TensorFlow 使用示例和最佳实践

import tensorflow as tf
import numpy as np

# 创建张量的多种方式
def tensor_creation_examples():
    """张量创建示例
    
    功能说明:
    - 演示各种张量创建方法
    - 展示数据类型和形状控制
    - 介绍设备放置策略
    """
    
    # 1. 从常量创建
    constant_tensor = tf.constant([1, 2, 3, 4], dtype=tf.float32)
    print(f"常量张量: {constant_tensor}")
    
    # 2. 从numpy数组创建
    numpy_array = np.array([[1, 2], [3, 4]])
    tensor_from_numpy = tf.constant(numpy_array)
    print(f"从numpy创建: {tensor_from_numpy}")
    
    # 3. 创建特殊张量
    zeros_tensor = tf.zeros((3, 3))          # 零张量
    ones_tensor = tf.ones((2, 4))            # 全1张量
    random_tensor = tf.random.normal((2, 3)) # 随机张量
    
    # 4. 变量张量（可训练）
    variable_tensor = tf.Variable(
        initial_value=tf.random.normal((3, 3)),
        trainable=True,
        name="my_variable"
    )
    
    # 5. 指定设备
    with tf.device('/CPU:0'):
        cpu_tensor = tf.constant([1, 2, 3])
    
    # 如果有GPU可用
    if tf.config.list_physical_devices('GPU'):
        with tf.device('/GPU:0'):
            gpu_tensor = tf.constant([4, 5, 6])
    
    return {
        'constant': constant_tensor,
        'from_numpy': tensor_from_numpy,
        'zeros': zeros_tensor,
        'ones': ones_tensor,
        'random': random_tensor,
        'variable': variable_tensor
    }

# 张量运算示例
def tensor_operations_examples():
    """张量运算示例
    
    功能说明:
    - 演示基础数学运算
    - 展示广播机制
    - 介绍形状操作
    """
    
    # 创建示例张量
    a = tf.constant([[1, 2], [3, 4]], dtype=tf.float32)
    b = tf.constant([[5, 6], [7, 8]], dtype=tf.float32)
    
    # 基础运算
    addition = tf.add(a, b)              # 加法
    multiplication = tf.multiply(a, b)    # 逐元素乘法
    matrix_mult = tf.matmul(a, b)        # 矩阵乘法
    
    # 使用运算符重载
    addition_op = a + b
    multiplication_op = a * b
    matrix_mult_op = a @ b
    
    # 广播运算
    scalar = tf.constant(2.0)
    broadcast_mult = a * scalar
    
    # 形状操作
    reshaped = tf.reshape(a, (4, 1))     # 重塑形状
    transposed = tf.transpose(a)          # 转置
    expanded = tf.expand_dims(a, axis=0)  # 增加维度
    
    # 聚合操作
    sum_all = tf.reduce_sum(a)           # 所有元素求和
    sum_axis = tf.reduce_sum(a, axis=1)  # 按轴求和
    mean_val = tf.reduce_mean(a)         # 平均值
    max_val = tf.reduce_max(a)           # 最大值
    
    return {
        'addition': addition,
        'multiplication': multiplication,
        'matrix_mult': matrix_mult,
        'broadcast_mult': broadcast_mult,
        'reshaped': reshaped,
        'transposed': transposed,
        'sum_all': sum_all,
        'mean_val': mean_val
    }

# 使用Keras Sequential API
def build_sequential_model():
    """构建序列模型示例
    
    功能说明:
    - 演示Sequential API的使用
    - 展示不同类型的层
    - 介绍模型编译和训练
    """
    
    # 创建序列模型
    model = tf.keras.Sequential([
        # 输入层（可选，用于明确指定输入形状）
        tf.keras.layers.InputLayer(input_shape=(28, 28, 1)),
        
        # 卷积层
        tf.keras.layers.Conv2D(
            filters=32,
            kernel_size=(3, 3),
            activation='relu',
            padding='same'
        ),
        tf.keras.layers.MaxPooling2D(pool_size=(2, 2)),
        
        # 更多卷积层
        tf.keras.layers.Conv2D(64, (3, 3), activation='relu'),
        tf.keras.layers.MaxPooling2D((2, 2)),
        
        # 展平层
        tf.keras.layers.Flatten(),
        
        # 全连接层
        tf.keras.layers.Dense(128, activation='relu'),
        tf.keras.layers.Dropout(0.5),  # 防止过拟合
        
        # 输出层
        tf.keras.layers.Dense(10, activation='softmax')
    ])
    
    # 编译模型
    model.compile(
        optimizer='adam',
        loss='sparse_categorical_crossentropy',
        metrics=['accuracy']
    )
    
    # 打印模型结构
    model.summary()
    
    return model

# 使用Functional API
def build_functional_model():
    """构建函数式模型示例
    
    功能说明:
    - 演示Functional API的灵活性
    - 展示多输入多输出模型
    - 介绍模型的复杂连接
    """
    
    # 定义输入
    input_layer = tf.keras.layers.Input(shape=(28, 28, 1), name='input_image')
    
    # 卷积分支
    conv1 = tf.keras.layers.Conv2D(32, (3, 3), activation='relu')(input_layer)
    pool1 = tf.keras.layers.MaxPooling2D((2, 2))(conv1)
    
    conv2 = tf.keras.layers.Conv2D(64, (3, 3), activation='relu')(pool1)
    pool2 = tf.keras.layers.MaxPooling2D((2, 2))(conv2)
    
    # 展平
    flatten = tf.keras.layers.Flatten()(pool2)
    
    # 全连接分支
    dense1 = tf.keras.layers.Dense(128, activation='relu')(flatten)
    dropout1 = tf.keras.layers.Dropout(0.5)(dense1)
    
    # 多个输出
    output1 = tf.keras.layers.Dense(10, activation='softmax', name='classification')(dropout1)
    output2 = tf.keras.layers.Dense(1, activation='sigmoid', name='confidence')(dropout1)
    
    # 创建模型
    model = tf.keras.Model(
        inputs=input_layer,
        outputs=[output1, output2],
        name='multi_output_model'
    )
    
    # 编译模型（多输出需要指定多个损失函数）
    model.compile(
        optimizer='adam',
        loss={
            'classification': 'sparse_categorical_crossentropy',
            'confidence': 'binary_crossentropy'
        },
        metrics={
            'classification': ['accuracy'],
            'confidence': ['mae']
        }
    )
    
    return model

# 自定义层
class CustomDenseLayer(tf.keras.layers.Layer):
    """自定义全连接层示例
    
    功能说明:
    - 演示如何创建自定义层
    - 展示权重初始化和管理
    - 介绍前向传播的实现
    """
    
    def __init__(self, units, activation=None, **kwargs):
        """初始化自定义层
        
        参数:
            units: 输出单元数
            activation: 激活函数
        """
        super(CustomDenseLayer, self).__init__(**kwargs)
        self.units = units
        self.activation = tf.keras.activations.get(activation)
    
    def build(self, input_shape):
        """构建层，创建权重
        
        参数:
            input_shape: 输入形状
        """
        # 创建权重矩阵
        self.kernel = self.add_weight(
            name='kernel',
            shape=(input_shape[-1], self.units),
            initializer='glorot_uniform',
            trainable=True
        )
        
        # 创建偏置向量
        self.bias = self.add_weight(
            name='bias',
            shape=(self.units,),
            initializer='zeros',
            trainable=True
        )
        
        super(CustomDenseLayer, self).build(input_shape)
    
    def call(self, inputs, **kwargs):
        """前向传播
        
        参数:
            inputs: 输入张量
            
        返回:
            输出张量
        """
        # 线性变换: y = xW + b
        output = tf.matmul(inputs, self.kernel) + self.bias
        
        # 应用激活函数
        if self.activation is not None:
            output = self.activation(output)
        
        return output
    
    def get_config(self):
        """获取层配置，用于序列化"""
        config = super(CustomDenseLayer, self).get_config()
        config.update({
            'units': self.units,
            'activation': tf.keras.activations.serialize(self.activation)
        })
        return config

def complete_training_example():
    """完整的模型训练示例
    
    功能说明:
    - 演示数据准备和预处理
    - 展示训练循环的实现
    - 介绍模型评估和保存
    """
    
    # 1. 数据准备
    (x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()
    
    # 数据预处理
    x_train = x_train.astype('float32') / 255.0  # 归一化
    x_test = x_test.astype('float32') / 255.0
    
    # 添加通道维度
    x_train = tf.expand_dims(x_train, -1)
    x_test = tf.expand_dims(x_test, -1)
    
    # 2. 创建数据集
    train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train))
    train_dataset = train_dataset.shuffle(buffer_size=10000)
    train_dataset = train_dataset.batch(32)
    train_dataset = train_dataset.prefetch(tf.data.AUTOTUNE)
    
    test_dataset = tf.data.Dataset.from_tensor_slices((x_test, y_test))
    test_dataset = test_dataset.batch(32)
    
    # 3. 构建模型
    model = build_sequential_model()
    
    # 4. 设置回调函数
    callbacks = [
        # 早停
        tf.keras.callbacks.EarlyStopping(
            monitor='val_loss',
            patience=3,
            restore_best_weights=True
        ),
        
        # 学习率调度
        tf.keras.callbacks.ReduceLROnPlateau(
            monitor='val_loss',
            factor=0.2,
            patience=2,
            min_lr=1e-7
        ),
        
        # 模型检查点
        tf.keras.callbacks.ModelCheckpoint(
            filepath='model_checkpoint.h5',
            monitor='val_accuracy',
            save_best_only=True,
            save_weights_only=False
        ),
        
        # TensorBoard日志
        tf.keras.callbacks.TensorBoard(
            log_dir='./logs',
            histogram_freq=1,
            write_graph=True,
            write_images=True
        )
    ]
    
    # 5. 训练模型
    history = model.fit(
        train_dataset,
        epochs=10,
        validation_data=test_dataset,
        callbacks=callbacks,
        verbose=1
    )
    
    # 6. 评估模型
    test_loss, test_accuracy = model.evaluate(test_dataset, verbose=0)
    print(f"测试准确率: {test_accuracy:.4f}")
    
    # 7. 保存模型
    model.save('trained_model.h5')
    
    # 8. 可视化训练历史
    import matplotlib.pyplot as plt
    
    plt.figure(figsize=(12, 4))
    
    plt.subplot(1, 2, 1)
    plt.plot(history.history['loss'], label='训练损失')
    plt.plot(history.history['val_loss'], label='验证损失')
    plt.title('模型损失')
    plt.xlabel('轮次')
    plt.ylabel('损失')
    plt.legend()
    
    plt.subplot(1, 2, 2)
    plt.plot(history.history['accuracy'], label='训练准确率')
    plt.plot(history.history['val_accuracy'], label='验证准确率')
    plt.title('模型准确率')
    plt.xlabel('轮次')
    plt.ylabel('准确率')
    plt.legend()
    
    plt.tight_layout()
    plt.savefig('training_history.png')
    plt.show()
    
    return model, history

# tf.function基础用法
@tf.function
def basic_function_example(x, y):
    """tf.function基础示例
    
    功能说明:
    - 将Python函数转换为TensorFlow图
    - 提供更好的性能和可移植性
    - 支持自动微分和优化
    """
    z = tf.matmul(x, y)
    return tf.nn.relu(z)

# 带输入签名的tf.function
@tf.function(input_signature=[
    tf.TensorSpec(shape=[None, 784], dtype=tf.float32),
    tf.TensorSpec(shape=[784, 10], dtype=tf.float32)
])
def function_with_signature(x, w):
    """带输入签名的tf.function
    
    功能说明:
    - 明确指定输入的形状和类型
    - 避免不必要的重新跟踪
    - 提高函数调用效率
    """
    return tf.nn.softmax(tf.matmul(x, w))

# 条件控制流
@tf.function
def conditional_function(x, training=True):
    """条件控制流示例
    
    功能说明:
    - 在图中使用条件语句
    - 支持动态控制流
    - 处理训练和推理的不同逻辑
    """
    if training:
        # 训练时添加噪声
        noise = tf.random.normal(tf.shape(x), stddev=0.1)
        x = x + noise
    
    # 应用dropout（仅在训练时）
    x = tf.nn.dropout(x, rate=0.2 if training else 0.0)
    
    return x

# 循环控制流
@tf.function
def loop_function(x, n_iterations):
    """循环控制流示例
    
    功能说明:
    - 在图中使用循环
    - 支持动态循环次数
    - 优化循环性能
    """
    for i in tf.range(n_iterations):
        x = tf.nn.relu(x)
        x = x * 0.9  # 逐渐衰减
    
    return x

# 使用tf.while_loop进行更复杂的循环
@tf.function
def while_loop_function(x):
    """while循环示例
    
    功能说明:
    - 使用tf.while_loop进行条件循环
    - 支持复杂的循环逻辑
    - 提供更好的性能优化
    """
    def condition(i, x):
        return tf.reduce_mean(x) > 0.1
    
    def body(i, x):
        x = x * 0.9
        return i + 1, x
    
    final_i, final_x = tf.while_loop(
        condition, body, [0, x],
        shape_invariants=[tf.TensorShape([]), x.get_shape()]
    )
    
    return final_x

def gradient_computation_examples():
    """梯度计算示例
    
    功能说明:
    - 演示tf.GradientTape的使用
    - 展示高阶梯度计算
    - 介绍梯度的应用场景
    """
    
    # 基础梯度计算
    def basic_gradient():
        x = tf.Variable(3.0)
        
        with tf.GradientTape() as tape:
            y = x ** 2 + 2 * x + 1
        
        # 计算dy/dx
        gradient = tape.gradient(y, x)
        print(f"x = {x.numpy()}, y = {y.numpy()}, dy/dx = {gradient.numpy()}")
        
        return gradient
    
    # 多变量梯度
    def multi_variable_gradient():
        x = tf.Variable(2.0)
        y = tf.Variable(3.0)
        
        with tf.GradientTape() as tape:
            z = x**2 + y**2 + 2*x*y
        
        # 计算关于所有变量的梯度
        gradients = tape.gradient(z, [x, y])
        print(f"dz/dx = {gradients[0].numpy()}, dz/dy = {gradients[1].numpy()}")
        
        return gradients
    
    # 高阶梯度
    def higher_order_gradient():
        x = tf.Variable(2.0)
        
        with tf.GradientTape() as tape2:
            with tf.GradientTape() as tape1:
                y = x**3
            
            # 一阶梯度
            first_grad = tape1.gradient(y, x)
        
        # 二阶梯度
        second_grad = tape2.gradient(first_grad, x)
        print(f"d²y/dx² = {second_grad.numpy()}")
        
        return second_grad
    
    # 在神经网络中的应用
    def neural_network_gradient():
        # 创建简单的线性模型
        model = tf.keras.Sequential([
            tf.keras.layers.Dense(1, input_shape=(1,))
        ])
        
        # 准备数据
        x_data = tf.constant([[1.0], [2.0], [3.0], [4.0]])
        y_true = tf.constant([[2.0], [4.0], [6.0], [8.0]])
        
        # 定义损失函数
        def loss_fn(y_true, y_pred):
            return tf.reduce_mean(tf.square(y_true - y_pred))
        
        # 计算梯度
        with tf.GradientTape() as tape:
            y_pred = model(x_data)
            loss = loss_fn(y_true, y_pred)
        
        # 获取模型参数的梯度
        gradients = tape.gradient(loss, model.trainable_variables)
        
        print(f"损失: {loss.numpy()}")
        for i, grad in enumerate(gradients):
            print(f"参数 {i} 的梯度: {grad.numpy()}")
        
        return gradients
    
    # 执行所有示例
    basic_gradient()
    multi_variable_gradient()
    higher_order_gradient()
    neural_network_gradient()

def optimized_data_pipeline():
    """优化的数据管道示例
    
    功能说明:
    - 演示高效的数据加载和预处理
    - 展示并行化技术
    - 介绍缓存和预取策略
    """
    
    # 创建示例数据
    def create_sample_data():
        # 模拟图像数据
        images = tf.random.normal((1000, 224, 224, 3))
        labels = tf.random.uniform((1000,), maxval=10, dtype=tf.int32)
        return images, labels
    
    images, labels = create_sample_data()
    
    # 基础数据集
    dataset = tf.data.Dataset.from_tensor_slices((images, labels))
    
    # 数据预处理函数
    def preprocess_image(image, label):
        """图像预处理函数
        
        功能说明:
        - 图像归一化
        - 数据增强
        - 类型转换
        """
        # 归一化到[0, 1]
        image = tf.cast(image, tf.float32) / 255.0
        
        # 随机翻转（数据增强）
        image = tf.image.random_flip_left_right(image)
        
        # 随机亮度调整
        image = tf.image.random_brightness(image, max_delta=0.1)
        
        return image, label
    
    # 优化的数据管道
    optimized_dataset = (dataset
        .shuffle(buffer_size=1000)           # 打乱数据
        .map(preprocess_image,               # 并行预处理
             num_parallel_calls=tf.data.AUTOTUNE)
        .batch(32)                           # 批处理
        .prefetch(tf.data.AUTOTUNE)         # 预取数据
        .cache()                            # 缓存处理后的数据
    )
    
    # 进一步优化：使用tf.data.experimental
    advanced_dataset = (dataset
        .shuffle(buffer_size=1000)
        .map(preprocess_image, 
             num_parallel_calls=tf.data.AUTOTUNE)
        .batch(32)
        .prefetch(tf.data.AUTOTUNE)
        # 可选：将数据缓存到磁盘
        # .cache('/tmp/dataset_cache')
    )
    
    return optimized_dataset

def data_loading_best_practices():
    """数据加载最佳实践
    
    功能说明:
    - 展示不同数据源的处理方法
    - 介绍内存和性能优化技巧
    - 提供错误处理策略
    """
    
    # 1. 从文件加载数据
    def load_from_files():
        # 获取文件路径列表
        file_paths = tf.data.Dataset.list_files('/path/to/images/*.jpg')
        
        def parse_image(file_path):
            # 读取图像文件
            image = tf.io.read_file(file_path)
            image = tf.image.decode_jpeg(image, channels=3)
            image = tf.image.resize(image, [224, 224])
            return image
        
        dataset = file_paths.map(parse_image, 
                                num_parallel_calls=tf.data.AUTOTUNE)
        return dataset
    
    # 2. 处理大型数据集
    def handle_large_dataset():
        # 使用tf.data.TFRecordDataset处理大型数据集
        filenames = ['/path/to/train.tfrecord']
        dataset = tf.data.TFRecordDataset(filenames)
        
        def parse_tfrecord(example_proto):
            # 定义特征描述
            feature_description = {
                'image': tf.io.FixedLenFeature([], tf.string),
                'label': tf.io.FixedLenFeature([], tf.int64),
            }
            
            # 解析示例
            features = tf.io.parse_single_example(example_proto, feature_description)
            
            # 解码图像
            image = tf.io.decode_raw(features['image'], tf.uint8)
            image = tf.reshape(image, [224, 224, 3])
            
            label = tf.cast(features['label'], tf.int32)
            
            return image, label
        
        parsed_dataset = dataset.map(parse_tfrecord)
        return parsed_dataset
    
    # 3. 错误处理和恢复
    def error_handling_dataset():
        def potentially_failing_map(x):
            # 模拟可能失败的操作
            return tf.cond(
                tf.random.uniform(()) < 0.1,  # 10%的概率失败
                lambda: tf.constant(-1),       # 失败时返回-1
                lambda: x * 2                  # 成功时返回2x
            )
        
        dataset = tf.data.Dataset.range(100)
        
        # 使用ignore_errors()忽略错误的元素
        robust_dataset = (dataset
            .map(potentially_failing_map)
            .filter(lambda x: x >= 0)  # 过滤掉错误值
        )
        
        return robust_dataset

def model_optimization_techniques():
    """模型优化技术示例
    
    功能说明:
    - 演示混合精度训练
    - 展示模型量化技术
    - 介绍图优化方法
    """
    
    # 1. 混合精度训练
    def mixed_precision_training():
        """混合精度训练设置
        
        功能说明:
        - 使用float16进行前向传播
        - 使用float32进行梯度计算
        - 减少内存使用，提高训练速度
        """
        # 设置混合精度策略
        policy = tf.keras.mixed_precision.Policy('mixed_float16')
        tf.keras.mixed_precision.set_global_policy(policy)
        
        # 创建模型
        model = tf.keras.Sequential([
            tf.keras.layers.Dense(128, activation='relu'),
            tf.keras.layers.Dense(64, activation='relu'),
            # 输出层必须使用float32
            tf.keras.layers.Dense(10, activation='softmax', dtype='float32')
        ])
        
        # 编译模型
        optimizer = tf.keras.optimizers.Adam()
        # 包装优化器以处理混合精度
        optimizer = tf.keras.mixed_precision.LossScaleOptimizer(optimizer)
        
        model.compile(
            optimizer=optimizer,
            loss='sparse_categorical_crossentropy',
            metrics=['accuracy']
        )
        
        return model
    
    # 2. 模型量化
    def model_quantization():
        """模型量化示例
        
        功能说明:
        - 将float32权重量化为int8
        - 减少模型大小和推理时间
        - 保持相对较高的精度
        """
        # 创建并训练一个简单模型
        model = tf.keras.Sequential([
            tf.keras.layers.Dense(128, activation='relu', input_shape=(784,)),
            tf.keras.layers.Dense(10, activation='softmax')
        ])
        
        model.compile(optimizer='adam',
                     loss='sparse_categorical_crossentropy',
                     metrics=['accuracy'])
        
        # 准备代表性数据集用于量化
        def representative_dataset():
            for _ in range(100):
                yield [tf.random.normal((1, 784))]
        
        # 转换为TensorFlow Lite模型并量化
        converter = tf.lite.TFLiteConverter.from_keras_model(model)
        converter.optimizations = [tf.lite.Optimize.DEFAULT]
        converter.representative_dataset = representative_dataset
        converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
        converter.inference_input_type = tf.int8
        converter.inference_output_type = tf.int8
        
        quantized_model = converter.convert()
        
        # 保存量化模型
        with open('quantized_model.tflite', 'wb') as f:
            f.write(quantized_model)
        
        return quantized_model
    
    # 3. 图优化
    @tf.function
    def optimized_inference_function(model, x):
        """优化的推理函数
        
        功能说明:
        - 使用tf.function进行图优化
        - 减少Python开销
        - 启用各种图级优化
        """
        return model(x, training=False)
    
    # 4. 内存优化
    def memory_optimization():
        """内存优化技术
        
        功能说明:
        - 梯度检查点
        - 内存增长控制
        - 显存管理
        """
        # 配置GPU内存增长
        gpus = tf.config.experimental.list_physical_devices('GPU')
        if gpus:
            try:
                for gpu in gpus:
                    tf.config.experimental.set_memory_growth(gpu, True)
            except RuntimeError as e:
                print(e)
        
        # 使用梯度检查点减少内存使用
        class MemoryEfficientModel(tf.keras.Model):
            def __init__(self):
                super().__init__()
                self.dense_layers = [
                    tf.keras.layers.Dense(1024, activation='relu')
                    for _ in range(10)
                ]
                self.output_layer = tf.keras.layers.Dense(10)
            
            @tf.recompute_grad
            def call(self, inputs):
                x = inputs
                for layer in self.dense_layers:
                    x = layer(x)
                return self.output_layer(x)
        
        return MemoryEfficientModel()

def custom_python_operations():
    """Python自定义操作示例
    
    功能说明:
    - 演示如何创建自定义Python操作
    - 展示tf.py_function的使用
    - 介绍梯度定义方法
    """
    
    # 1. 使用tf.py_function
    def custom_numpy_operation(x, y):
        """使用numpy实现的自定义操作"""
        import numpy as np
        
        # 自定义的numpy操作
        result = np.sin(x.numpy()) * np.cos(y.numpy())
        return tf.constant(result, dtype=tf.float32)
    
    @tf.function
    def use_custom_operation():
        x = tf.constant([1.0, 2.0, 3.0])
        y = tf.constant([4.0, 5.0, 6.0])
        
        # 使用tf.py_function调用Python函数
        result = tf.py_function(
            func=custom_numpy_operation,
            inp=[x, y],
            Tout=tf.float32
        )
        
        # 设置输出形状（tf.py_function无法推断形状）
        result.set_shape(x.shape)
        
        return result
    
    # 2. 自定义梯度
    @tf.custom_gradient
    def custom_relu(x):
        """带自定义梯度的ReLU函数
        
        功能说明:
        - 实现自定义的前向传播
        - 定义对应的反向传播
        - 支持自动微分
        """
        # 前向传播
        result = tf.nn.relu(x)
        
        def grad_fn(dy):
            """自定义梯度函数
            
            参数:
                dy: 上游梯度
                
            返回:
                关于输入的梯度
            """
            # ReLU的梯度：x > 0时为1，否则为0
            return dy * tf.cast(x > 0, tf.float32)
        
        return result, grad_fn
    
    # 3. 复杂的自定义操作
    @tf.custom_gradient
    def gumbel_softmax(logits, temperature=1.0):
        """Gumbel Softmax操作
        
        功能说明:
        - 实现可微分的离散采样
        - 用于强化学习和变分推断
        - 提供温度参数控制
        """
        # 添加Gumbel噪声
        gumbel_noise = -tf.math.log(-tf.math.log(
            tf.random.uniform(tf.shape(logits))
        ))
        
        # 计算Gumbel Softmax
        y = tf.nn.softmax((logits + gumbel_noise) / temperature)
        
        def grad_fn(dy):
            """Gumbel Softmax的梯度"""
            # 计算softmax的雅可比矩阵
            s = tf.nn.softmax(logits / temperature)
            jacobian = (s * dy - s * tf.reduce_sum(s * dy, axis=-1, keepdims=True)) / temperature
            return jacobian, None  # temperature不需要梯度
        
        return y, grad_fn
    
    # 测试自定义操作
    def test_custom_operations():
        # 测试自定义ReLU
        x = tf.constant([-1.0, 0.0, 1.0, 2.0])
        
        with tf.GradientTape() as tape:
            tape.watch(x)
            y = custom_relu(x)
        
        gradients = tape.gradient(y, x)
        print(f"自定义ReLU输出: {y.numpy()}")
        print(f"自定义ReLU梯度: {gradients.numpy()}")
        
        # 测试Gumbel Softmax
        logits = tf.constant([[1.0, 2.0, 3.0], [2.0, 1.0, 0.0]])
        
        with tf.GradientTape() as tape:
            tape.watch(logits)
            samples = gumbel_softmax(logits, temperature=0.5)
        
        gradients = tape.gradient(samples, logits)
        print(f"Gumbel Softmax样本: {samples.numpy()}")
        print(f"Gumbel Softmax梯度形状: {gradients.shape}")
    
    test_custom_operations()
    return use_custom_operation()

// custom_ops.cc - C++自定义操作示例
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/shape_inference.h"

using namespace tensorflow;

// 注册自定义操作
REGISTER_OP("CustomMatMul")
    .Input("a: T")
    .Input("b: T")
    .Output("product: T")
    .Attr("T: {float, double}")
    .Attr("transpose_a: bool = false")
    .Attr("transpose_b: bool = false")
    .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
        // 形状推断逻辑
        shape_inference::ShapeHandle a_shape, b_shape;
        TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 2, &a_shape));
        TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 2, &b_shape));
        
        // 获取维度
        shape_inference::DimensionHandle a_rows = c->Dim(a_shape, 0);
        shape_inference::DimensionHandle a_cols = c->Dim(a_shape, 1);
        shape_inference::DimensionHandle b_rows = c->Dim(b_shape, 0);
        shape_inference::DimensionHandle b_cols = c->Dim(b_shape, 1);
        
        // 检查矩阵乘法的维度兼容性
        shape_inference::DimensionHandle merged;
        TF_RETURN_IF_ERROR(c->Merge(a_cols, b_rows, &merged));
        
        // 设置输出形状
        c->set_output(0, c->Matrix(a_rows, b_cols));
        return absl::OkStatus();
    })
    .Doc(R"doc(
自定义矩阵乘法操作

该操作执行两个矩阵的乘法运算，支持转置选项。

a: 第一个输入矩阵
b: 第二个输入矩阵
product: 矩阵乘法的结果
transpose_a: 是否转置矩阵a
transpose_b: 是否转置矩阵b
)doc");

// 实现CPU版本的操作内核
template <typename T>
class CustomMatMulOp : public OpKernel {
public:
    explicit CustomMatMulOp(OpKernelConstruction* context) : OpKernel(context) {
        // 获取属性
        OP_REQUIRES_OK(context, context->GetAttr("transpose_a", &transpose_a_));
        OP_REQUIRES_OK(context, context->GetAttr("transpose_b", &transpose_b_));
    }
    
    void Compute(OpKernelContext* context) override {
        // 获取输入张量
        const Tensor& a = context->input(0);
        const Tensor& b = context->input(1);
        
        // 验证输入
        OP_REQUIRES(context, a.dims() == 2,
                    errors::InvalidArgument("输入a必须是2维矩阵"));
        OP_REQUIRES(context, b.dims() == 2,
                    errors::InvalidArgument("输入b必须是2维矩阵"));
        
        // 获取矩阵维度
        int64_t a_rows = transpose_a_ ? a.dim_size(1) : a.dim_size(0);
        int64_t a_cols = transpose_a_ ? a.dim_size(0) : a.dim_size(1);
        int64_t b_rows = transpose_b_ ? b.dim_size(1) : b.dim_size(0);
        int64_t b_cols = transpose_b_ ? b.dim_size(0) : b.dim_size(1);
        
        // 检查维度兼容性
        OP_REQUIRES(context, a_cols == b_rows,
                    errors::InvalidArgument("矩阵维度不兼容"));
        
        // 分配输出张量
        Tensor* output = nullptr;
        OP_REQUIRES_OK(context, context->allocate_output(
            0, TensorShape({a_rows, b_cols}), &output));
        
        // 执行矩阵乘法
        auto a_matrix = a.matrix<T>();
        auto b_matrix = b.matrix<T>();
        auto output_matrix = output->matrix<T>();
        
        // 简单的矩阵乘法实现（实际应用中应使用优化的BLAS库）
        for (int64_t i = 0; i < a_rows; ++i) {
            for (int64_t j = 0; j < b_cols; ++j) {
                T sum = T(0);
                for (int64_t k = 0; k < a_cols; ++k) {
                    T a_val = transpose_a_ ? a_matrix(k, i) : a_matrix(i, k);
                    T b_val = transpose_b_ ? b_matrix(j, k) : b_matrix(k, j);
                    sum += a_val * b_val;
                }
                output_matrix(i, j) = sum;
            }
        }
    }

private:
    bool transpose_a_;
    bool transpose_b_;
};

// 注册CPU内核
REGISTER_KERNEL_BUILDER(
    Name("CustomMatMul").Device(DEVICE_CPU).TypeConstraint<float>("T"),
    CustomMatMulOp<float>);

REGISTER_KERNEL_BUILDER(
    Name("CustomMatMul").Device(DEVICE_CPU).TypeConstraint<double>("T"),
    CustomMatMulOp<double>);

# 使用自定义C++操作的Python包装
def load_custom_ops():
    """加载自定义C++操作
    
    功能说明:
    - 编译和加载自定义操作库
    - 提供Python接口
    - 处理错误和异常
    """
    
    # 加载自定义操作库
    import tensorflow as tf
    
    # 假设已经编译了自定义操作库
    custom_ops_module = tf.load_op_library('./custom_ops.so')
    
    def custom_matmul(a, b, transpose_a=False, transpose_b=False, name=None):
        """自定义矩阵乘法的Python包装
        
        参数:
            a: 第一个输入矩阵
            b: 第二个输入矩阵
            transpose_a: 是否转置矩阵a
            transpose_b: 是否转置矩阵b
            name: 操作名称
            
        返回:
            矩阵乘法结果
        """
        return custom_ops_module.custom_mat_mul(
            a, b, transpose_a=transpose_a, transpose_b=transpose_b, name=name
        )
    
    # 测试自定义操作
    def test_custom_matmul():
        a = tf.constant([[1.0, 2.0], [3.0, 4.0]])
        b = tf.constant([[5.0, 6.0], [7.0, 8.0]])
        
        # 使用自定义操作
        result_custom = custom_matmul(a, b)
        
        # 使用标准操作进行比较
        result_standard = tf.matmul(a, b)
        
        print(f"自定义操作结果:\n{result_custom.numpy()}")
        print(f"标准操作结果:\n{result_standard.numpy()}")
        print(f"结果是否相等: {tf.reduce_all(tf.equal(result_custom, result_standard)).numpy()}")
    
    test_custom_matmul()
    return custom_matmul