Article

TensorFlow 模板

更新于:2026-07-17

通用训练模板

import tensorflow as tf
from tensorflow.keras import layers, models, optimizers, callbacks, losses, metrics
import os
import logging
import time
import numpy as np


# ================== Step 1: 数据准备 ==================
def get_dataset(batch_size=32, input_shape=(32, 32, 3), num_classes=10):
    """
    示例数据:随机生成模拟数据(可替换为真实数据集)
    """
    x_train = np.random.rand(1000, *input_shape).astype(np.float32)
    y_train = np.random.randint(0, num_classes, (1000,))
    x_val = np.random.rand(200, *input_shape).astype(np.float32)
    y_val = np.random.randint(0, num_classes, (200,))

    train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train))
    val_dataset = tf.data.Dataset.from_tensor_slices((x_val, y_val))

    # 增强与预处理(可自定义)
    def preprocess(x, y):
        return x, y

    train_dataset = train_dataset.map(preprocess).shuffle(1000).batch(batch_size).prefetch(tf.data.AUTOTUNE)
    val_dataset = val_dataset.map(preprocess).batch(batch_size).prefetch(tf.data.AUTOTUNE)

    return train_dataset, val_dataset


# ================== Step 2: 构建模型 ==================
def build_model(input_shape=(32, 32, 3), num_classes=10):
    """
    示例模型:一个简单的 CNN 模型,可替换为你自己的网络结构
    """
    model = models.Sequential([
        layers.Input(shape=input_shape),
        layers.Rescaling(1./255),  # 可选归一化层
        layers.Conv2D(32, 3, activation='relu'),
        layers.MaxPooling2D(),
        layers.Conv2D(64, 3, activation='relu'),
        layers.MaxPooling2D(),
        layers.Flatten(),
        layers.Dense(64, activation='relu'),
        layers.Dense(num_classes)
    ])

    return model


# ================== Step 3: 工具函数与配置 ==================
def setup_logger(log_dir="logs"):
    os.makedirs(log_dir, exist_ok=True)
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(levelname)s - %(message)s',
        handlers=[
            logging.FileHandler(f"{log_dir}/train_{time.strftime('%Y%m%d_%H%M%S')}.log"),
            logging.StreamHandler()
        ]
    )


# ================== Step 4: 主程序入口 ==================
def main():
    setup_logger()
    strategy = tf.distribute.MirroredStrategy()
    logging.info(f'Number of devices: {strategy.num_replicas_in_sync}')

    with strategy.scope():
        batch_size = 64
        epochs = 10
        input_shape = (32, 32, 3)
        num_classes = 10

        # 获取数据集(替换为你自己的 Dataset)
        train_dataset, val_dataset = get_dataset(batch_size, input_shape, num_classes)

        # 构建模型(可替换为 ResNet、EfficientNet 等)
        model = build_model(input_shape, num_classes)

        # 编译模型
        optimizer = optimizers.Adam(learning_rate=1e-3)
        loss_fn = losses.SparseCategoricalCrossentropy(from_logits=True)
        metric = metrics.SparseCategoricalAccuracy(name='acc')

        model.compile(optimizer=optimizer,
                      loss=loss_fn,
                      metrics=[metric])

        # 回调函数
        checkpoint_cb = callbacks.ModelCheckpoint(
            "checkpoints/best_model.keras",
            monitor='val_loss',
            save_best_only=True,
            mode='min'
        )
        early_stop_cb = callbacks.EarlyStopping(patience=5, restore_best_weights=True)
        tensorboard_cb = callbacks.TensorBoard(log_dir="runs")

        callbacks_list = [checkpoint_cb, early_stop_cb, tensorboard_cb]

        # 开始训练
        history = model.fit(
            train_dataset,
            validation_data=val_dataset,
            epochs=epochs,
            callbacks=callbacks_list
        )

        logging.info("Training completed.")
        model.save("checkpoints/final_model.keras")


if __name__ == "__main__":
    main()

输出说明

文件夹内容
logs/存放训练日志文件
runs/TensorBoard 日志,用于可视化训练过程
checkpoints/保存最佳模型和最终模型

扩展建议

你可以基于此模板轻松添加以下功能:

功能描述
使用预训练模型替换 build_model() 中的网络为 ResNet50, EfficientNet, MobileNetV2 等
图像增强preprocess() 函数中加入 layers.RandomFlip(), layers.RandomRotation()
多 GPU 训练使用 tf.distribute.MirroredStrategy() 支持多卡并行
混合精度训练添加 policy = mixed_precision.Policy('mixed_float16') 提升训练效率
自动学习率调度添加 ReduceLROnPlateau 或使用 CosineDecay 调度器
自定义损失函数替换 SparseCategoricalCrossentropyFocalLossDiceLoss
模型评估与推理加载 .keras 模型后进行预测、混淆矩阵分析等
使用 TFRecords将大型数据集转换为 TFRecord 格式提高加载效率

总结

这个 TensorFlow 模板具有以下优点:

特性描述
简洁清晰结构分明,适合初学者和项目启动
模块化设计各组件解耦,便于更换和复用
支持 TensorBoard实时监控训练过程
支持多 GPU利用 MirroredStrategy 加速训练
易于扩展可快速接入各种新功能
支持混合精度提高训练速度和内存效率

LeNet

下面是一个基于 TensorFlow 泛用训练模板,实现经典卷积神经网络模型 LeNet-5 的完整代码示例。我们将使用 MNIST 手写数字数据集 进行训练和验证。

使用 TensorFlow 通用模板实现 LeNet

import tensorflow as tf
from tensorflow.keras import layers, models, optimizers, callbacks, losses, metrics
import os
import logging
import time


# ================== Step 1: 数据准备 ==================
def get_dataset(batch_size=32):
    """
    加载并预处理 MNIST 数据集
    """
    (x_train, y_train), (x_val, y_val) = tf.keras.datasets.mnist.load_data()

    # 增加通道维度并归一化
    x_train = x_train[..., tf.newaxis].astype("float32") / 255.0
    x_val = x_val[..., tf.newaxis].astype("float32") / 255.0

    train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train))
    val_dataset = tf.data.Dataset.from_tensor_slices((x_val, y_val))

    # 数据增强(可选)
    def preprocess(x, y):
        return x, y

    train_dataset = train_dataset.map(preprocess).shuffle(10000).batch(batch_size).prefetch(tf.data.AUTOTUNE)
    val_dataset = val_dataset.map(preprocess).batch(batch_size).prefetch(tf.data.AUTOTUNE)

    return train_dataset, val_dataset


# ================== Step 2: 构建 LeNet 模型 ==================
def build_model(num_classes=10):
    """
    构建 LeNet-5 模型结构
    输入尺寸:(28, 28, 1)
    输出类别数:num_classes
    """
    model = models.Sequential([
        layers.Input(shape=(28, 28, 1)),

        # Layer 1: Conv + AvgPool
        layers.Conv2D(6, kernel_size=5, activation='tanh'),
        layers.AveragePooling2D(pool_size=2),

        # Layer 2: Conv + AvgPool
        layers.Conv2D(16, kernel_size=5, activation='tanh'),
        layers.AveragePooling2D(pool_size=2),

        # Flatten
        layers.Flatten(),

        # Fully Connected Layers
        layers.Dense(120, activation='tanh'),
        layers.Dense(84, activation='tanh'),
        layers.Dense(num_classes)
    ])

    return model


# ================== Step 3: 工具函数与配置 ==================
def setup_logger(log_dir="logs"):
    os.makedirs(log_dir, exist_ok=True)
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(levelname)s - %(message)s',
        handlers=[
            logging.FileHandler(f"{log_dir}/train_{time.strftime('%Y%m%d_%H%M%S')}.log"),
            logging.StreamHandler()
        ]
    )


# ================== Step 4: 主程序入口 ==================
def main():
    setup_logger()
    strategy = tf.distribute.MirroredStrategy()
    logging.info(f'Number of devices: {strategy.num_replicas_in_sync}')

    with strategy.scope():
        batch_size = 128
        epochs = 10
        num_classes = 10

        # 获取数据集
        train_dataset, val_dataset = get_dataset(batch_size)

        # 构建模型
        model = build_model(num_classes)

        # 编译模型
        optimizer = optimizers.Adam(learning_rate=1e-3)
        loss_fn = losses.SparseCategoricalCrossentropy(from_logits=True)
        metric = metrics.SparseCategoricalAccuracy(name='acc')

        model.compile(optimizer=optimizer,
                      loss=loss_fn,
                      metrics=[metric])

        # 回调函数
        checkpoint_cb = callbacks.ModelCheckpoint(
            "checkpoints/best_lenet.keras",
            monitor='val_acc',
            save_best_only=True,
            mode='max'
        )
        early_stop_cb = callbacks.EarlyStopping(patience=5, restore_best_weights=True)
        tensorboard_cb = callbacks.TensorBoard(log_dir="runs")

        callbacks_list = [checkpoint_cb, early_stop_cb, tensorboard_cb]

        # 开始训练
        history = model.fit(
            train_dataset,
            validation_data=val_dataset,
            epochs=epochs,
            callbacks=callbacks_list
        )

        logging.info("Training completed.")
        model.save("checkpoints/final_lenet.keras")


if __name__ == "__main__":
    main()

输出说明

文件夹内容
logs/存放训练日志文件
runs/TensorBoard 日志,用于可视化训练过程
checkpoints/保存最佳模型 best_lenet.keras 和最终模型 final_lenet.keras

关键说明

模型结构说明(LeNet-5)

层级类型参数说明
Conv2D(6, 5×5)卷积层6 个 5×5 卷积核,tanh 激活
AvgPool(2×2)平均池化下采样到 14×14
Conv2D(16, 5×5)卷积层16 个 5×5 卷积核,tanh 激活
AvgPool(2×2)平均池化下采样到 5×5
Dense(120)全连接层120 神经元
Dense(84)全连接层84 神经元
Dense(10)输出层10 分类输出(logits)

数据预处理

  • 使用了标准的 MNIST 数据集。
  • 图像大小为 (28, 28, 1),添加了通道维度。
  • 使用了 tf.data.Dataset 来高效加载数据,并支持 shuffle、batch 和 prefetch。

模型编译

  • 使用 Adam 优化器,学习率默认为 1e-3
  • 损失函数为 SparseCategoricalCrossentropy(适用于整数标签)
  • 指标为准确率 SparseCategoricalAccuracy

扩展建议

功能描述
使用 Sigmoid 替换 Tanh更现代的做法,可尝试不同激活函数效果
添加 Dropout 层防止过拟合
使用 BatchNormalization提高收敛速度和稳定性
使用 TensorBoard 可视化在浏览器中查看损失和准确率曲线
多 GPU 训练使用 MirroredStrategy 提升训练效率
自定义评估指标如 Precision、Recall、F1-score
模型推理加载模型后进行预测或部署

VGG

下面是一个基于 TensorFlow 通用训练模板实现的经典深度学习模型 —— VGG16 的完整代码示例。我们将使用 CIFAR-10 数据集 来进行训练和验证。由于 VGG16 是为 ImageNet 设计的(输入大小为 224×224),为了适应 CIFAR-10(32×32),我们对网络结构进行了适当的调整。

使用 TensorFlow 通用模板实现 VGG16

import tensorflow as tf
from tensorflow.keras import layers, models, optimizers, callbacks, losses, metrics
import os
import logging
import time


# ================== Step 1: 数据准备 ==================
def get_dataset(batch_size=64):
    """
    加载并预处理 CIFAR-10 数据集
    """
    (x_train, y_train), (x_val, y_val) = tf.keras.datasets.cifar10.load_data()

    # 归一化 + 类型转换
    x_train = x_train.astype("float32") / 255.0
    x_val = x_val.astype("float32") / 255.0

    train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train))
    val_dataset = tf.data.Dataset.from_tensor_slices((x_val, y_val))

    # 数据增强(可选)
    def preprocess(x, y):
        x = tf.image.random_flip_left_right(x)
        x = tf.image.random_brightness(x, max_delta=0.2)
        return x, tf.squeeze(y)

    train_dataset = train_dataset.map(preprocess).shuffle(10000).batch(batch_size).prefetch(tf.data.AUTOTUNE)
    val_dataset = val_dataset.map(lambda x, y: (x, tf.squeeze(y))).batch(batch_size).prefetch(tf.data.AUTOTUNE)

    return train_dataset, val_dataset


# ================== Step 2: 构建 VGG16 模型 ==================
def build_model(num_classes=10):
    """
    构建简化版 VGG16 模型(适配 CIFAR-10 输入尺寸 32x32x3)
    """
    model = models.Sequential([
        layers.Input(shape=(32, 32, 3)),

        # Block 1
        layers.Conv2D(64, kernel_size=3, padding='same', activation='relu'),
        layers.Conv2D(64, kernel_size=3, padding='same', activation='relu'),
        layers.MaxPooling2D(pool_size=2, strides=2),

        # Block 2
        layers.Conv2D(128, kernel_size=3, padding='same', activation='relu'),
        layers.Conv2D(128, kernel_size=3, padding='same', activation='relu'),
        layers.MaxPooling2D(pool_size=2, strides=2),

        # Block 3
        layers.Conv2D(256, kernel_size=3, padding='same', activation='relu'),
        layers.Conv2D(256, kernel_size=3, padding='same', activation='relu'),
        layers.Conv2D(256, kernel_size=3, padding='same', activation='relu'),
        layers.MaxPooling2D(pool_size=2, strides=2),

        # Block 4
        layers.Conv2D(512, kernel_size=3, padding='same', activation='relu'),
        layers.Conv2D(512, kernel_size=3, padding='same', activation='relu'),
        layers.Conv2D(512, kernel_size=3, padding='same', activation='relu'),
        layers.MaxPooling2D(pool_size=2, strides=2),

        # Block 5
        layers.Conv2D(512, kernel_size=3, padding='same', activation='relu'),
        layers.Conv2D(512, kernel_size=3, padding='same', activation='relu'),
        layers.Conv2D(512, kernel_size=3, padding='same', activation='relu'),
        layers.MaxPooling2D(pool_size=2, strides=2),

        # Flatten
        layers.Flatten(),

        # Fully Connected Layers
        layers.Dense(4096, activation='relu'),
        layers.Dropout(0.5),
        layers.Dense(4096, activation='relu'),
        layers.Dropout(0.5),
        layers.Dense(num_classes)
    ])

    return model


# ================== Step 3: 工具函数与配置 ==================
def setup_logger(log_dir="logs"):
    os.makedirs(log_dir, exist_ok=True)
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(levelname)s - %(message)s',
        handlers=[
            logging.FileHandler(f"{log_dir}/train_{time.strftime('%Y%m%d_%H%M%S')}.log"),
            logging.StreamHandler()
        ]
    )


# ================== Step 4: 主程序入口 ==================
def main():
    setup_logger()
    strategy = tf.distribute.MirroredStrategy()
    logging.info(f'Number of devices: {strategy.num_replicas_in_sync}')

    with strategy.scope():
        batch_size = 64
        epochs = 50
        num_classes = 10

        # 获取数据集
        train_dataset, val_dataset = get_dataset(batch_size)

        # 构建模型
        model = build_model(num_classes)

        # 编译模型
        optimizer = optimizers.Adam(learning_rate=1e-4)
        loss_fn = losses.SparseCategoricalCrossentropy(from_logits=True)
        metric = metrics.SparseCategoricalAccuracy(name='acc')

        model.compile(optimizer=optimizer,
                      loss=loss_fn,
                      metrics=[metric])

        # 回调函数
        checkpoint_cb = callbacks.ModelCheckpoint(
            "checkpoints/best_vgg16.keras",
            monitor='val_acc',
            save_best_only=True,
            mode='max'
        )
        early_stop_cb = callbacks.EarlyStopping(patience=10, restore_best_weights=True)
        reduce_lr = callbacks.ReduceLROnPlateau(monitor='val_loss', factor=0.2, patience=5)
        tensorboard_cb = callbacks.TensorBoard(log_dir="runs")

        callbacks_list = [checkpoint_cb, early_stop_cb, reduce_lr, tensorboard_cb]

        # 开始训练
        history = model.fit(
            train_dataset,
            validation_data=val_dataset,
            epochs=epochs,
            callbacks=callbacks_list
        )

        logging.info("Training completed.")
        model.save("checkpoints/final_vgg16.keras")


if __name__ == "__main__":
    main()

输出说明

文件夹内容
logs/存放训练日志文件
runs/TensorBoard 日志,用于可视化训练过程
checkpoints/保存最佳模型 best_vgg16.keras 和最终模型 final_vgg16.keras

关键说明

模型结构说明(简化版 VGG16)

层级类型参数说明
Conv2D(64, 3×3)卷积层64 个卷积核,ReLU 激活
MaxPool池化层下采样
Conv2D(128, 3×3)卷积层更深层特征提取
多个卷积层提取高维语义信息
Dense(4096)全连接层高维表示
Dropout正则化防止过拟合
Dense(10)输出层分类输出(logits)

数据预处理

  • 使用了标准的 CIFAR-10 数据集。
  • 图像大小为 (32, 32, 3),已归一化到 [0, 1] 范围。
  • 使用了 tf.data.Dataset 进行高效加载和批处理,并应用了基本的数据增强技术(如随机水平翻转和亮度调整)。

模型编译

  • 使用 Adam 优化器,学习率默认为 1e-4
  • 损失函数为 SparseCategoricalCrossentropy(适用于整数标签)
  • 指标为准确率 SparseCategoricalAccuracy

训练策略

  • 使用了 ModelCheckpoint 保存最佳模型
  • 使用了 ReduceLROnPlateau 动态调整学习率
  • 使用了 EarlyStopping 防止过拟合
  • 使用了 TensorBoard 可视化训练过程

扩展建议

功能描述
增加图像增强添加 RandomFlipRandomRotation 等数据增强操作
替换为原始 VGG16将输入改为 224×224,适配 ImageNet
使用 BatchNormalization提升收敛速度和泛化能力
使用预训练模型如 Keras Applications 中的 ResNet、EfficientNet 等
多 GPU 训练使用 MirroredStrategy 加速训练
自定义损失函数如 Focal Loss、Label Smoothing
模型推理加载模型后进行预测或部署

GoogLeNet / Inception

下面是一个基于 TensorFlow 通用训练模板实现的经典深度学习模型 —— GoogLeNet(Inception v1) 的完整代码示例。我们将使用 CIFAR-10 数据集 来进行训练和验证。由于 GoogLeNet 是为 ImageNet 设计的(输入大小为 224×224),为了适应 CIFAR-10(32×32),我们对网络结构进行了适当的调整。

使用 TensorFlow 通用模板实现 GoogLeNet (Inception v1)

import tensorflow as tf
from tensorflow.keras import layers, models, optimizers, callbacks, losses, metrics
import os
import logging
import time


# ================== Step 1: 数据准备 ==================
def get_dataset(batch_size=64):
    """
    加载并预处理 CIFAR-10 数据集
    """
    (x_train, y_train), (x_val, y_val) = tf.keras.datasets.cifar10.load_data()

    # 归一化 + 类型转换
    x_train = x_train.astype("float32") / 255.0
    x_val = x_val.astype("float32") / 255.0

    train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train))
    val_dataset = tf.data.Dataset.from_tensor_slices((x_val, y_val))

    # 数据增强(可选)
    def preprocess(x, y):
        x = tf.image.random_flip_left_right(x)
        x = tf.image.random_brightness(x, max_delta=0.2)
        return x, tf.squeeze(y)

    train_dataset = train_dataset.map(preprocess).shuffle(10000).batch(batch_size).prefetch(tf.data.AUTOTUNE)
    val_dataset = val_dataset.map(lambda x, y: (x, tf.squeeze(y))).batch(batch_size).prefetch(tf.data.AUTOTUNE)

    return train_dataset, val_dataset


# ================== Step 2: 构建 Inception 模块 ==================
class InceptionModule(layers.Layer):
    def __init__(self, filters_branches, **kwargs):
        super(InceptionModule, self).__init__(**kwargs)
        filter1, filter3_reduce, filter3, filter5_reduce, filter5, pool_proj = filters_branches

        self.branch1 = layers.Conv2D(filter1, kernel_size=1, activation='relu', padding='same')
        self.branch2 = tf.keras.Sequential([
            layers.Conv2D(filter3_reduce, kernel_size=1, activation='relu', padding='same'),
            layers.Conv2D(filter3, kernel_size=3, activation='relu', padding='same')
        ])
        self.branch3 = tf.keras.Sequential([
            layers.Conv2D(filter5_reduce, kernel_size=1, activation='relu', padding='same'),
            layers.Conv2D(filter5, kernel_size=5, activation='relu', padding='same')
        ])
        self.branch4 = tf.keras.Sequential([
            layers.MaxPooling2D(pool_size=3, strides=1, padding='same'),
            layers.Conv2D(pool_proj, kernel_size=1, activation='relu', padding='same')
        ])

    def call(self, inputs, **kwargs):
        branch1 = self.branch1(inputs)
        branch2 = self.branch2(inputs)
        branch3 = self.branch3(inputs)
        branch4 = self.branch4(inputs)
        outputs = layers.concatenate([branch1, branch2, branch3, branch4])
        return outputs


# ================== Step 3: 构建 GoogLeNet 模型 ==================
def build_model(num_classes=10):
    """
    构建简化版 GoogLeNet 模型(适配 CIFAR-10 输入尺寸 32x32x3)
    """
    model = models.Sequential([
        layers.Input(shape=(32, 32, 3)),

        # Initial Convolutional Layer
        layers.Conv2D(64, kernel_size=7, strides=2, padding='same', activation='relu'),
        layers.MaxPooling2D(pool_size=3, strides=2, padding='same'),

        # Inception Module 1
        InceptionModule(filters_branches=[64, 96, 128, 16, 32, 32]),
        InceptionModule(filters_branches=[128, 128, 192, 32, 96, 64]),
        layers.MaxPooling2D(pool_size=3, strides=2, padding='same'),

        # Inception Module 2
        InceptionModule(filters_branches=[192, 96, 208, 16, 48, 64]),
        InceptionModule(filters_branches=[160, 112, 224, 24, 64, 64]),

        # Flatten
        layers.Flatten(),

        # Fully Connected Layers
        layers.Dense(1024, activation='relu'),
        layers.Dropout(0.4),
        layers.Dense(num_classes)
    ])

    return model


# ================== Step 4: 工具函数与配置 ==================
def setup_logger(log_dir="logs"):
    os.makedirs(log_dir, exist_ok=True)
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(levelname)s - %(message)s',
        handlers=[
            logging.FileHandler(f"{log_dir}/train_{time.strftime('%Y%m%d_%H%M%S')}.log"),
            logging.StreamHandler()
        ]
    )


# ================== Step 5: 主程序入口 ==================
def main():
    setup_logger()
    strategy = tf.distribute.MirroredStrategy()
    logging.info(f'Number of devices: {strategy.num_replicas_in_sync}')

    with strategy.scope():
        batch_size = 64
        epochs = 50
        num_classes = 10

        # 获取数据集
        train_dataset, val_dataset = get_dataset(batch_size)

        # 构建模型
        model = build_model(num_classes)

        # 编译模型
        optimizer = optimizers.Adam(learning_rate=1e-4)
        loss_fn = losses.SparseCategoricalCrossentropy(from_logits=True)
        metric = metrics.SparseCategoricalAccuracy(name='acc')

        model.compile(optimizer=optimizer,
                      loss=loss_fn,
                      metrics=[metric])

        # 回调函数
        checkpoint_cb = callbacks.ModelCheckpoint(
            "checkpoints/best_googlenet.keras",
            monitor='val_acc',
            save_best_only=True,
            mode='max'
        )
        early_stop_cb = callbacks.EarlyStopping(patience=10, restore_best_weights=True)
        reduce_lr = callbacks.ReduceLROnPlateau(monitor='val_loss', factor=0.2, patience=5)
        tensorboard_cb = callbacks.TensorBoard(log_dir="runs")

        callbacks_list = [checkpoint_cb, early_stop_cb, reduce_lr, tensorboard_cb]

        # 开始训练
        history = model.fit(
            train_dataset,
            validation_data=val_dataset,
            epochs=epochs,
            callbacks=callbacks_list
        )

        logging.info("Training completed.")
        model.save("checkpoints/final_googlenet.keras")


if __name__ == "__main__":
    main()

输出说明

文件夹内容
logs/存放训练日志文件
runs/TensorBoard 日志,用于可视化训练过程
checkpoints/保存最佳模型 best_googlenet.keras 和最终模型 final_googlenet.keras

关键说明

模型结构说明(简化版 GoogLeNet)

  • Inception 模块:通过并行使用不同大小的卷积核(1×1、3×3、5×5)和最大池化层来提取特征,然后将这些结果拼接起来形成输出。
  • Initial Convolutional Layer:初始卷积层,用于初步特征提取。
  • Fully Connected Layers:全连接层,用于分类任务。

数据预处理

  • 使用了标准的 CIFAR-10 数据集。
  • 图像大小为 (32, 32, 3),已归一化到 [0, 1] 范围。
  • 使用了 tf.data.Dataset 进行高效加载和批处理,并应用了基本的数据增强技术(如随机水平翻转和亮度调整)。

模型编译

  • 使用 Adam 优化器,学习率默认为 1e-4
  • 损失函数为 SparseCategoricalCrossentropy(适用于整数标签)
  • 指标为准确率 SparseCategoricalAccuracy

训练策略

  • 使用了 ModelCheckpoint 保存最佳模型
  • 使用了 ReduceLROnPlateau 动态调整学习率
  • 使用了 EarlyStopping 防止过拟合
  • 使用了 TensorBoard 可视化训练过程

扩展建议

功能描述
增加图像增强添加 RandomFlipRandomRotation 等数据增强操作
替换为原始 GoogLeNet将输入改为 224×224,适配 ImageNet
使用 BatchNormalization提升收敛速度和泛化能力
多 GPU 训练使用 MirroredStrategy 加速训练
自定义损失函数如 Focal Loss、Label Smoothing
模型推理加载模型后进行预测或部署

ResNet

下面是一个基于 TensorFlow 通用训练模板实现的经典深度学习模型 —— ResNet-34 的完整代码示例。我们将使用 CIFAR-10 数据集 来进行训练和验证。由于 ResNet 是为 ImageNet 设计的(输入大小为 224×224),为了适应 CIFAR-10(32×32),我们对网络结构进行了适当的调整。

使用 TensorFlow 通用模板实现 ResNet-34

import tensorflow as tf
from tensorflow.keras import layers, models, optimizers, callbacks, losses, metrics
import os
import logging
import time


# ================== Step 1: 数据准备 ==================
def get_dataset(batch_size=64):
    """
    加载并预处理 CIFAR-10 数据集
    """
    (x_train, y_train), (x_val, y_val) = tf.keras.datasets.cifar10.load_data()

    # 归一化 + 类型转换
    x_train = x_train.astype("float32") / 255.0
    x_val = x_val.astype("float32") / 255.0

    train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train))
    val_dataset = tf.data.Dataset.from_tensor_slices((x_val, y_val))

    # 数据增强(可选)
    def preprocess(x, y):
        x = tf.image.random_flip_left_right(x)
        x = tf.image.random_brightness(x, max_delta=0.2)
        return x, tf.squeeze(y)

    train_dataset = train_dataset.map(preprocess).shuffle(10000).batch(batch_size).prefetch(tf.data.AUTOTUNE)
    val_dataset = val_dataset.map(lambda x, y: (x, tf.squeeze(y))).batch(batch_size).prefetch(tf.data.AUTOTUNE)

    return train_dataset, val_dataset


# ================== Step 2: 构建残差块 ==================
class ResidualBlock(layers.Layer):
    def __init__(self, filters, strides=1, use_projection=False, **kwargs):
        super(ResidualBlock, self).__init__(**kwargs)
        self.use_projection = use_projection

        self.conv1 = layers.Conv2D(filters, kernel_size=3, strides=strides, padding='same')
        self.bn1 = layers.BatchNormalization()
        self.relu1 = layers.ReLU()

        self.conv2 = layers.Conv2D(filters, kernel_size=3, padding='same')
        self.bn2 = layers.BatchNormalization()

        if self.use_projection:
            self.shortcut = layers.Conv2D(filters, kernel_size=1, strides=strides, padding='same')

        self.add = layers.Add()
        self.relu_out = layers.ReLU()

    def call(self, inputs, training=None):
        shortcut = inputs
        if self.use_projection:
            shortcut = self.shortcut(shortcut)

        x = self.conv1(inputs)
        x = self.bn1(x, training=training)
        x = self.relu1(x)

        x = self.conv2(x)
        x = self.bn2(x, training=training)

        x = self.add([x, shortcut])
        return self.relu_out(x)


# ================== Step 3: 构建 ResNet-34 模型 ==================
def build_model(num_classes=10):
    """
    构建简化版 ResNet-34 模型(适配 CIFAR-10 输入尺寸 32x32x3)
    """
    inputs = layers.Input(shape=(32, 32, 3))

    # Initial Convolutional Layer
    x = layers.Conv2D(64, kernel_size=7, strides=2, padding='same')(inputs)
    x = layers.BatchNormalization()(x)
    x = layers.ReLU()(x)
    x = layers.MaxPooling2D(pool_size=3, strides=2, padding='same')(x)

    # Residual Blocks
    def residual_block_sequence(filters, num_blocks, strides_first_block=1):
        x = ResidualBlock(filters=filters, strides=strides_first_block, use_projection=True)(x)
        for _ in range(1, num_blocks):
            x = ResidualBlock(filters=filters)(x)
        return x

    x = residual_block_sequence(64, num_blocks=3, strides_first_block=1)
    x = residual_block_sequence(128, num_blocks=4, strides_first_block=2)
    x = residual_block_sequence(256, num_blocks=6, strides_first_block=2)
    x = residual_block_sequence(512, num_blocks=3, strides_first_block=2)

    # Final Layers
    x = layers.GlobalAveragePooling2D()(x)
    outputs = layers.Dense(num_classes)(x)

    model = models.Model(inputs=inputs, outputs=outputs)
    return model


# ================== Step 4: 工具函数与配置 ==================
def setup_logger(log_dir="logs"):
    os.makedirs(log_dir, exist_ok=True)
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(levelname)s - %(message)s',
        handlers=[
            logging.FileHandler(f"{log_dir}/train_{time.strftime('%Y%m%d_%H%M%S')}.log"),
            logging.StreamHandler()
        ]
    )


# ================== Step 5: 主程序入口 ==================
def main():
    setup_logger()
    strategy = tf.distribute.MirroredStrategy()
    logging.info(f'Number of devices: {strategy.num_replicas_in_sync}')

    with strategy.scope():
        batch_size = 64
        epochs = 50
        num_classes = 10

        # 获取数据集
        train_dataset, val_dataset = get_dataset(batch_size)

        # 构建模型
        model = build_model(num_classes)

        # 编译模型
        optimizer = optimizers.Adam(learning_rate=1e-4)
        loss_fn = losses.SparseCategoricalCrossentropy(from_logits=True)
        metric = metrics.SparseCategoricalAccuracy(name='acc')

        model.compile(optimizer=optimizer,
                      loss=loss_fn,
                      metrics=[metric])

        # 回调函数
        checkpoint_cb = callbacks.ModelCheckpoint(
            "checkpoints/best_resnet34.keras",
            monitor='val_acc',
            save_best_only=True,
            mode='max'
        )
        early_stop_cb = callbacks.EarlyStopping(patience=10, restore_best_weights=True)
        reduce_lr = callbacks.ReduceLROnPlateau(monitor='val_loss', factor=0.2, patience=5)
        tensorboard_cb = callbacks.TensorBoard(log_dir="runs")

        callbacks_list = [checkpoint_cb, early_stop_cb, reduce_lr, tensorboard_cb]

        # 开始训练
        history = model.fit(
            train_dataset,
            validation_data=val_dataset,
            epochs=epochs,
            callbacks=callbacks_list
        )

        logging.info("Training completed.")
        model.save("checkpoints/final_resnet34.keras")


if __name__ == "__main__":
    main()

输出说明

文件夹内容
logs/存放训练日志文件
runs/TensorBoard 日志,用于可视化训练过程
checkpoints/保存最佳模型 best_resnet34.keras 和最终模型 final_resnet34.keras

关键说明

模型结构说明(简化版 ResNet-34)

  • 残差块(Residual Block):通过跳过连接(skip connections)将输入直接加到卷积层输出上,解决了深层网络中的梯度消失问题。
  • Initial Convolutional Layer:初始卷积层,用于初步特征提取。
  • Final Layers:全局平均池化层和全连接层,用于分类任务。

数据预处理

  • 使用了标准的 CIFAR-10 数据集。
  • 图像大小为 (32, 32, 3),已归一化到 [0, 1] 范围。
  • 使用了 tf.data.Dataset 进行高效加载和批处理,并应用了基本的数据增强技术(如随机水平翻转和亮度调整)。

模型编译

  • 使用 Adam 优化器,学习率默认为 1e-4
  • 损失函数为 SparseCategoricalCrossentropy(适用于整数标签)
  • 指标为准确率 SparseCategoricalAccuracy

训练策略

  • 使用了 ModelCheckpoint 保存最佳模型
  • 使用了 ReduceLROnPlateau 动态调整学习率
  • 使用了 EarlyStopping 防止过拟合
  • 使用了 TensorBoard 可视化训练过程

扩展建议

功能描述
增加图像增强添加 RandomFlipRandomRotation 等数据增强操作
替换为原始 ResNet将输入改为 224×224,适配 ImageNet
使用 BatchNormalization提升收敛速度和泛化能力
多 GPU 训练使用 MirroredStrategy 加速训练
自定义损失函数如 Focal Loss、Label Smoothing
模型推理加载模型后进行预测或部署

U-Net

U-Net 是一种常用于医学图像分割的卷积神经网络模型。它特别适合于小数据集训练,通过使用数据增强、加权损失函数和高效的下采样与上采样路径来实现精确的像素级分类。

下面是一个基于 TensorFlow 通用训练模板实现 U-Net 模型,并以一个假设的数据集为例进行训练的完整代码示例。我们将不具体指定数据集,但会提供一个框架,您可以根据自己的数据集(如医学影像分割任务)调整输入部分。

使用 TensorFlow 通用模板实现 U-Net

import tensorflow as tf
from tensorflow.keras import layers, models, optimizers, callbacks, losses
import numpy as np
import os
import logging
import time


# ================== Step 1: 数据准备(这里以随机生成的数据为例)==================
def generate_random_dataset(num_samples=1000, img_shape=(128, 128, 1), mask_shape=(128, 128, 1)):
    """
    随机生成数据集供演示使用。
    """
    images = np.random.rand(num_samples, *img_shape)
    masks = np.random.randint(0, 2, size=(num_samples, *mask_shape))
    return images, masks

def get_dataset(batch_size=32):
    """
    加载并预处理数据集
    """
    images, masks = generate_random_dataset()

    dataset = tf.data.Dataset.from_tensor_slices((images, masks))
    dataset = dataset.shuffle(buffer_size=1024).batch(batch_size).prefetch(tf.data.AUTOTUNE)

    # 分割为训练集和验证集
    train_size = int(0.8 * len(images))
    train_dataset = dataset.take(train_size)
    val_dataset = dataset.skip(train_size)

    return train_dataset, val_dataset


# ================== Step 2: 构建 U-Net 模型 ==================
def conv_block(input_tensor, num_filters):
    x = layers.Conv2D(num_filters, (3, 3), padding='same')(input_tensor)
    x = layers.BatchNormalization()(x)
    x = layers.Activation('relu')(x)
    x = layers.Conv2D(num_filters, (3, 3), padding='same')(x)
    x = layers.BatchNormalization()(x)
    x = layers.Activation('relu')(x)
    return x

def encoder_block(input_tensor, num_filters):
    x = conv_block(input_tensor, num_filters)
    p = layers.MaxPooling2D((2, 2))(x)
    return x, p

def decoder_block(input_tensor, skip_features, num_filters):
    x = layers.Conv2DTranspose(num_filters, (2, 2), strides=2, padding='same')(input_tensor)
    x = layers.concatenate([x, skip_features])
    x = conv_block(x, num_filters)
    return x

def build_unet(img_shape):
    inputs = layers.Input(shape=img_shape)

    s1, p1 = encoder_block(inputs, 64)
    s2, p2 = encoder_block(p1, 128)
    s3, p3 = encoder_block(p2, 256)
    s4, p4 = encoder_block(p3, 512)

    b = conv_block(p4, 1024)

    d1 = decoder_block(b, s4, 512)
    d2 = decoder_block(d1, s3, 256)
    d3 = decoder_block(d2, s2, 128)
    d4 = decoder_block(d3, s1, 64)

    outputs = layers.Conv2D(1, (1, 1), activation='sigmoid')(d4)

    model = models.Model(inputs=[inputs], outputs=[outputs])
    return model


# ================== Step 3: 工具函数与配置 ==================
def setup_logger(log_dir="logs"):
    os.makedirs(log_dir, exist_ok=True)
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(levelname)s - %(message)s',
        handlers=[
            logging.FileHandler(f"{log_dir}/train_{time.strftime('%Y%m%d_%H%M%S')}.log"),
            logging.StreamHandler()
        ]
    )


# ================== Step 4: 主程序入口 ==================
def main():
    setup_logger()
    strategy = tf.distribute.MirroredStrategy()
    logging.info(f'Number of devices: {strategy.num_replicas_in_sync}')

    with strategy.scope():
        batch_size = 32
        epochs = 50
        img_shape = (128, 128, 1)

        # 获取数据集
        train_dataset, val_dataset = get_dataset(batch_size)

        # 构建模型
        model = build_unet(img_shape)

        # 编译模型
        optimizer = optimizers.Adam(learning_rate=1e-4)
        loss_fn = losses.BinaryCrossentropy(from_logits=False)
        metric = 'accuracy'

        model.compile(optimizer=optimizer,
                      loss=loss_fn,
                      metrics=[metric])

        # 回调函数
        checkpoint_cb = callbacks.ModelCheckpoint(
            "checkpoints/best_unet.keras",
            monitor='val_accuracy',
            save_best_only=True,
            mode='max'
        )
        early_stop_cb = callbacks.EarlyStopping(patience=10, restore_best_weights=True)
        reduce_lr = callbacks.ReduceLROnPlateau(monitor='val_loss', factor=0.2, patience=5)
        tensorboard_cb = callbacks.TensorBoard(log_dir="runs")

        callbacks_list = [checkpoint_cb, early_stop_cb, reduce_lr, tensorboard_cb]

        # 开始训练
        history = model.fit(
            train_dataset,
            validation_data=val_dataset,
            epochs=epochs,
            callbacks=callbacks_list
        )

        logging.info("Training completed.")
        model.save("checkpoints/final_unet.keras")


if __name__ == "__main__":
    main()

输出说明

文件夹内容
logs/存放训练日志文件
runs/TensorBoard 日志,用于可视化训练过程
checkpoints/保存最佳模型 best_unet.keras 和最终模型 final_unet.keras

关键说明

模型结构说明(U-Net)

  • 编码器(Encoder):通过一系列的卷积和最大池化操作提取特征图。每个块包括两个卷积层,接着是批量归一化和 ReLU 激活函数。
  • 解码器(Decoder):通过转置卷积操作逐步恢复空间分辨率。每个块接收来自编码器对应阶段的特征映射作为跳过连接,以帮助保留位置信息。
  • 输出层:应用一个卷积层将通道数减少到所需的类别数量(本例中为 1),并使用 Sigmoid 激活函数进行二分类。

数据预处理

  • 示例中使用了随机生成的数据集。在实际应用中,您需要替换这部分代码以加载和预处理您的特定数据集,例如医学影像数据。

模型编译

  • 使用 Adam 优化器,学习率默认为 1e-4
  • 损失函数为 BinaryCrossentropy(适用于二分类问题)
  • 指标为准确率

训练策略

  • 使用了 ModelCheckpoint 保存最佳模型
  • 使用了 ReduceLROnPlateau 动态调整学习率
  • 使用了 EarlyStopping 防止过拟合
  • 使用了 TensorBoard 可视化训练过程

扩展建议

功能描述
增加图像增强添加 RandomFlipRandomRotation 等数据增强操作
替换为真实数据集如医学影像数据集等
多 GPU 训练使用 MirroredStrategy 加速训练
自定义损失函数如 Dice Loss 或 Jaccard Index 等更适合分割任务的损失函数
模型推理加载模型后进行预测或部署

GANs (DCGAN)

下面是一个基于 TensorFlow 通用训练模板实现的经典 GAN(生成对抗网络) 的完整代码示例。我们将使用 DCGAN(深度卷积 GAN) 架构,在 MNIST 手写数字数据集 上进行训练,以演示如何从随机噪声生成逼真的图像。

使用 TensorFlow 通用模板实现 DCGAN

import tensorflow as tf
from tensorflow.keras import layers, models, optimizers, callbacks
from tensorflow.keras.datasets import mnist
import numpy as np
import os
import logging
import time
import matplotlib.pyplot as plt


# ================== Step 1: 数据准备 ==================
def get_dataset(batch_size=128):
    """
    加载并预处理 MNIST 数据集
    """
    (train_images, _), _ = mnist.load_data()
    train_images = train_images.reshape(-1, 28, 28, 1).astype("float32") / 127.5 - 1.0  # [-1, 1] 归一化

    dataset = tf.data.Dataset.from_tensor_slices(train_images)
    dataset = dataset.shuffle(buffer_size=60000).batch(batch_size).prefetch(tf.data.AUTOTUNE)

    return dataset


# ================== Step 2: 构建生成器和判别器 ==================

def build_generator(latent_dim):
    """
    生成器:将随机噪声向量映射为 28x28x1 图像
    """
    model = models.Sequential([
        layers.Dense(7 * 7 * 256, use_bias=False, input_shape=(latent_dim,)),
        layers.BatchNormalization(),
        layers.ReLU(),

        layers.Reshape((7, 7, 256)),

        layers.Conv2DTranspose(128, (5, 5), strides=(1, 1), padding='same', use_bias=False),
        layers.BatchNormalization(),
        layers.ReLU(),

        layers.Conv2DTranspose(64, (5, 5), strides=(2, 2), padding='same', use_bias=False),
        layers.BatchNormalization(),
        layers.ReLU(),

        layers.Conv2DTranspose(1, (5, 5), strides=(2, 2), padding='same', use_bias=False, activation='tanh')
    ])

    return model


def build_discriminator():
    """
    判别器:判断图像是真实的还是生成的
    """
    model = models.Sequential([
        layers.Input(shape=(28, 28, 1)),

        layers.Conv2D(64, (5, 5), strides=(2, 2), padding='same'),
        layers.LeakyReLU(alpha=0.2),

        layers.Conv2D(128, (5, 5), strides=(2, 2), padding='same'),
        layers.LeakyReLU(alpha=0.2),

        layers.Flatten(),
        layers.Dense(1)  # 输出一个实数分数(不加 sigmoid)
    ])

    return model


# ================== Step 3: 定义 GAN 模型 ==================
class DCGAN(models.Model):
    def __init__(self, generator, discriminator, latent_dim, **kwargs):
        super(DCGAN, self).__init__(**kwargs)
        self.generator = generator
        self.discriminator = discriminator
        self.latent_dim = latent_dim

    def compile(self, g_optimizer, d_optimizer, loss_fn):
        super(DCGAN, self).compile()
        self.g_optimizer = g_optimizer
        self.d_optimizer = d_optimizer
        self.loss_fn = loss_fn

    def train_step(self, real_images):

        batch_size = tf.shape(real_images)[0]

        # ----------------------------
        # 训练判别器
        # ----------------------------
        noise = tf.random.normal(shape=(batch_size, self.latent_dim))
        generated_images = self.generator(noise)

        with tf.GradientTape() as tape:
            real_output = self.discriminator(real_images, training=True)
            fake_output = self.discriminator(generated_images, training=True)

            real_loss = self.loss_fn(tf.ones_like(real_output), real_output)
            fake_loss = self.loss_fn(tf.zeros_like(fake_output), fake_output)
            d_loss = real_loss + fake_loss

        grads = tape.gradient(d_loss, self.discriminator.trainable_weights)
        self.d_optimizer.apply_gradients(zip(grads, self.discriminator.trainable_weights))

        # ----------------------------
        # 训练生成器
        # ----------------------------
        noise = tf.random.normal(shape=(batch_size, self.latent_dim))

        with tf.GradientTape() as tape:
            generated_images = self.generator(noise, training=True)
            fake_output = self.discriminator(generated_images, training=False)

            g_loss = self.loss_fn(tf.ones_like(fake_output), fake_output)

        grads = tape.gradient(g_loss, self.generator.trainable_weights)
        self.g_optimizer.apply_gradients(zip(grads, self.generator.trainable_weights))

        return {"d_loss": d_loss, "g_loss": g_loss}


# ================== Step 4: 工具函数与配置 ==================
def setup_logger(log_dir="logs"):
    os.makedirs(log_dir, exist_ok=True)
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(levelname)s - %(message)s',
        handlers=[
            logging.FileHandler(f"{log_dir}/train_{time.strftime('%Y%m%d_%H%M%S')}.log"),
            logging.StreamHandler()
        ]
    )

def plot_generated_images(epoch, generator, test_input, save_path="images"):
    predictions = generator(test_input, training=False)
    fig, axes = plt.subplots(4, 4, figsize=(8, 8))
    for i, ax in enumerate(axes.flat):
        ax.imshow(predictions[i, :, :, 0] * 0.5 + 0.5, cmap='gray')
        ax.axis("off")
    plt.tight_layout()
    os.makedirs(save_path, exist_ok=True)
    plt.savefig(f"{save_path}/image_at_epoch_{epoch:04d}.png")
    plt.close()


# ================== Step 5: 主程序入口 ==================
def main():
    setup_logger()
    strategy = tf.distribute.MirroredStrategy()
    logging.info(f'Number of devices: {strategy.num_replicas_in_sync}')

    with strategy.scope():
        batch_size = 128
        epochs = 100
        latent_dim = 100

        # 获取数据集
        dataset = get_dataset(batch_size)

        # 构建模型
        generator = build_generator(latent_dim)
        discriminator = build_discriminator()

        # 编译模型
        dcgan = DCGAN(generator, discriminator, latent_dim)
        dcgan.compile(
            g_optimizer=optimizers.Adam(learning_rate=2e-4, beta_1=0.5),
            d_optimizer=optimizers.Adam(learning_rate=2e-4, beta_1=0.5),
            loss_fn=losses.Hinge()  # 可选 Hinge Loss 或 BinaryCrossentropy
        )

        # 回调函数
        checkpoint_cb = callbacks.ModelCheckpoint(
            "checkpoints/generator.keras",
            monitor='g_loss',
            save_best_only=True,
            mode='min'
        )
        tensorboard_cb = callbacks.TensorBoard(log_dir="runs")

        fixed_noise = tf.random.normal([16, latent_dim])  # 固定噪声用于可视化生成效果

        # 自定义回调:定期保存生成图像
        class ImageSaveCallback(callbacks.Callback):
            def on_epoch_end(self, epoch, logs=None):
                if (epoch + 1) % 5 == 0:
                    plot_generated_images(epoch + 1, self.model.generator, fixed_noise)

        image_save_cb = ImageSaveCallback()

        callbacks_list = [checkpoint_cb, tensorboard_cb, image_save_cb]

        # 开始训练
        history = dcgan.fit(
            dataset,
            epochs=epochs,
            callbacks=callbacks_list
        )

        logging.info("Training completed.")
        dcgan.generator.save("checkpoints/final_generator.keras")
        dcgan.discriminator.save("checkpoints/final_discriminator.keras")


if __name__ == "__main__":
    main()

输出说明

文件夹内容
logs/存放训练日志文件
runs/TensorBoard 日志,用于可视化训练过程
images/每隔一定轮次保存的生成图像
checkpoints/保存最佳生成器、最终生成器和判别器

关键说明

模型结构说明(DCGAN)

部分类型参数说明
GeneratorConv2DTranspose将噪声向量映射为真实图像
DiscriminatorConv2D判断图像是真实还是生成的
Loss FunctionHinge Loss更适合稳定训练 GAN
OptimizerAdam(2e-4, beta=0.5)常用于 GAN 的优化器参数

数据预处理

  • 使用了标准的 MNIST 数据集
  • 图像大小为 (28, 28, 1),归一化到 [-1, 1]
  • 使用 tf.data.Dataset 进行高效加载和批处理

模型编译与训练

  • 使用自定义 train_step() 来实现交替训练策略
  • 判别器使用真实标签 [1] 和生成标签 [0] 来计算损失
  • 生成器的目标是让判别器输出接近 [1]

扩展建议

功能描述
替换为其他架构如 CycleGAN、StyleGAN、Wasserstein-GAN 等
多 GPU 训练使用 MirroredStrategy 加速训练
自定义损失函数如 Wasserstein loss、LS-GAN loss
图像转译任务如 pix2pix、CycleGAN 等条件 GAN
模型推理用固定噪声生成图像或图像修复等

AlexNet

下面是一个基于 TensorFlow 通用训练模板实现的经典深度学习模型 —— AlexNet 的完整代码示例。我们将使用 CIFAR-10 数据集 来进行训练和验证,因为原始 AlexNet 是为 ImageNet 设计的(输入大小为 224×224),但为了演示目的,我们将其适配到 CIFAR-10(32×32)。

使用 TensorFlow 通用模板实现 AlexNet

import tensorflow as tf
from tensorflow.keras import layers, models, optimizers, callbacks, losses, metrics
import os
import logging
import time


# ================== Step 1: 数据准备 ==================
def get_dataset(batch_size=64):
    """
    加载并预处理 CIFAR-10 数据集
    """
    (x_train, y_train), (x_val, y_val) = tf.keras.datasets.cifar10.load_data()

    # 归一化 + 类型转换
    x_train = x_train.astype("float32") / 255.0
    x_val = x_val.astype("float32") / 255.0

    y_train = tf.squeeze(y_train)
    y_val = tf.squeeze(y_val)

    train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train))
    val_dataset = tf.data.Dataset.from_tensor_slices((x_val, y_val))

    # 数据增强(可选)
    def preprocess(x, y):
        return x, y

    train_dataset = train_dataset.map(preprocess).shuffle(10000).batch(batch_size).prefetch(tf.data.AUTOTUNE)
    val_dataset = val_dataset.map(preprocess).batch(batch_size).prefetch(tf.data.AUTOTUNE)

    return train_dataset, val_dataset


# ================== Step 2: 构建 AlexNet 模型 ==================
def build_model(num_classes=10):
    """
    构建简化版 AlexNet 模型(适配 CIFAR-10 输入尺寸 32x32x3)
    """
    model = models.Sequential([
        layers.Input(shape=(32, 32, 3)),

        # Layer 1
        layers.Conv2D(96, kernel_size=3, strides=1, padding='same'),
        layers.Activation('relu'),
        layers.MaxPooling2D(pool_size=3, strides=2),

        # Layer 2
        layers.Conv2D(256, kernel_size=3, padding='same'),
        layers.Activation('relu'),
        layers.MaxPooling2D(pool_size=3, strides=2),

        # Layer 3
        layers.Conv2D(384, kernel_size=3, padding='same'),
        layers.Activation('relu'),

        # Layer 4
        layers.Conv2D(384, kernel_size=3, padding='same'),
        layers.Activation('relu'),

        # Layer 5
        layers.Conv2D(256, kernel_size=3, padding='same'),
        layers.Activation('relu'),
        layers.MaxPooling2D(pool_size=3, strides=2),

        # Flatten
        layers.Flatten(),

        # Fully Connected Layers
        layers.Dense(4096),
        layers.Activation('relu'),
        layers.Dropout(0.5),

        layers.Dense(4096),
        layers.Activation('relu'),
        layers.Dropout(0.5),

        layers.Dense(num_classes)
    ])

    return model


# ================== Step 3: 工具函数与配置 ==================
def setup_logger(log_dir="logs"):
    os.makedirs(log_dir, exist_ok=True)
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(levelname)s - %(message)s',
        handlers=[
            logging.FileHandler(f"{log_dir}/train_{time.strftime('%Y%m%d_%H%M%S')}.log"),
            logging.StreamHandler()
        ]
    )


# ================== Step 4: 主程序入口 ==================
def main():
    setup_logger()
    strategy = tf.distribute.MirroredStrategy()
    logging.info(f'Number of devices: {strategy.num_replicas_in_sync}')

    with strategy.scope():
        batch_size = 128
        epochs = 30
        num_classes = 10

        # 获取数据集
        train_dataset, val_dataset = get_dataset(batch_size)

        # 构建模型
        model = build_model(num_classes)

        # 编译模型
        optimizer = optimizers.Adam(learning_rate=1e-3)
        loss_fn = losses.SparseCategoricalCrossentropy(from_logits=True)
        metric = metrics.SparseCategoricalAccuracy(name='acc')

        model.compile(optimizer=optimizer,
                      loss=loss_fn,
                      metrics=[metric])

        # 回调函数
        checkpoint_cb = callbacks.ModelCheckpoint(
            "checkpoints/best_alexnet.keras",
            monitor='val_acc',
            save_best_only=True,
            mode='max'
        )
        early_stop_cb = callbacks.EarlyStopping(patience=7, restore_best_weights=True)
        reduce_lr = callbacks.ReduceLROnPlateau(monitor='val_loss', factor=0.2, patience=3)
        tensorboard_cb = callbacks.TensorBoard(log_dir="runs")

        callbacks_list = [checkpoint_cb, early_stop_cb, reduce_lr, tensorboard_cb]

        # 开始训练
        history = model.fit(
            train_dataset,
            validation_data=val_dataset,
            epochs=epochs,
            callbacks=callbacks_list
        )

        logging.info("Training completed.")
        model.save("checkpoints/final_alexnet.keras")


if __name__ == "__main__":
    main()

输出说明

文件夹内容
logs/存放训练日志文件
runs/TensorBoard 日志,用于可视化训练过程
checkpoints/保存最佳模型 best_alexnet.keras 和最终模型 final_alexnet.keras

关键说明

模型结构说明(简化版 AlexNet)

层级类型参数说明
Conv2D(96, 3×3)卷积层96 个卷积核,ReLU 激活
MaxPool池化层下采样
Conv2D(256, 3×3)卷积层更深层特征提取
多个卷积层提取高维语义信息
Dense(4096)全连接层高维表示
Dropout正则化防止过拟合
Dense(10)输出层分类输出(logits)

数据预处理

  • 使用了标准的 CIFAR-10 数据集。
  • 图像大小为 (32, 32, 3),已归一化到 [0, 1] 范围。
  • 使用了 tf.data.Dataset 进行高效加载和批处理。

模型编译

  • 使用 Adam 优化器,学习率默认为 1e-3
  • 损失函数为 SparseCategoricalCrossentropy(适用于整数标签)
  • 指标为准确率 SparseCategoricalAccuracy

训练策略

  • 使用了 ModelCheckpoint 保存最佳模型
  • 使用了 ReduceLROnPlateau 动态调整学习率
  • 使用了 EarlyStopping 防止过拟合
  • 使用了 TensorBoard 可视化训练过程

扩展建议

功能描述
增加图像增强添加 RandomFlipRandomRotation 等数据增强操作
替换为原始 AlexNet将输入改为 224×224,适配 ImageNet
使用 BatchNormalization提升收敛速度和泛化能力
使用预训练模型如 Keras Applications 中的 ResNet、EfficientNet 等
多 GPU 训练使用 MirroredStrategy 加速训练
自定义损失函数如 Focal Loss、Label Smoothing
模型推理加载模型后进行预测或部署

VGG

下面是一个基于 TensorFlow 通用训练模板实现的经典深度学习模型 —— VGG16 的完整代码示例。我们将使用 CIFAR-10 数据集 来进行训练和验证。由于 VGG16 是为 ImageNet 设计的(输入大小为 224×224),为了适应 CIFAR-10(32×32),我们对网络结构进行了适当的调整。

使用 TensorFlow 通用模板实现 VGG16

import tensorflow as tf
from tensorflow.keras import layers, models, optimizers, callbacks, losses, metrics
import os
import logging
import time


# ================== Step 1: 数据准备 ==================
def get_dataset(batch_size=64):
    """
    加载并预处理 CIFAR-10 数据集
    """
    (x_train, y_train), (x_val, y_val) = tf.keras.datasets.cifar10.load_data()

    # 归一化 + 类型转换
    x_train = x_train.astype("float32") / 255.0
    x_val = x_val.astype("float32") / 255.0

    train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train))
    val_dataset = tf.data.Dataset.from_tensor_slices((x_val, y_val))

    # 数据增强(可选)
    def preprocess(x, y):
        x = tf.image.random_flip_left_right(x)
        x = tf.image.random_brightness(x, max_delta=0.2)
        return x, tf.squeeze(y)

    train_dataset = train_dataset.map(preprocess).shuffle(10000).batch(batch_size).prefetch(tf.data.AUTOTUNE)
    val_dataset = val_dataset.map(lambda x, y: (x, tf.squeeze(y))).batch(batch_size).prefetch(tf.data.AUTOTUNE)

    return train_dataset, val_dataset


# ================== Step 2: 构建 VGG16 模型 ==================
def build_model(num_classes=10):
    """
    构建简化版 VGG16 模型(适配 CIFAR-10 输入尺寸 32x32x3)
    """
    model = models.Sequential([
        layers.Input(shape=(32, 32, 3)),

        # Block 1
        layers.Conv2D(64, kernel_size=3, padding='same', activation='relu'),
        layers.Conv2D(64, kernel_size=3, padding='same', activation='relu'),
        layers.MaxPooling2D(pool_size=2, strides=2),

        # Block 2
        layers.Conv2D(128, kernel_size=3, padding='same', activation='relu'),
        layers.Conv2D(128, kernel_size=3, padding='same', activation='relu'),
        layers.MaxPooling2D(pool_size=2, strides=2),

        # Block 3
        layers.Conv2D(256, kernel_size=3, padding='same', activation='relu'),
        layers.Conv2D(256, kernel_size=3, padding='same', activation='relu'),
        layers.Conv2D(256, kernel_size=3, padding='same', activation='relu'),
        layers.MaxPooling2D(pool_size=2, strides=2),

        # Block 4
        layers.Conv2D(512, kernel_size=3, padding='same', activation='relu'),
        layers.Conv2D(512, kernel_size=3, padding='same', activation='relu'),
        layers.Conv2D(512, kernel_size=3, padding='same', activation='relu'),
        layers.MaxPooling2D(pool_size=2, strides=2),

        # Block 5
        layers.Conv2D(512, kernel_size=3, padding='same', activation='relu'),
        layers.Conv2D(512, kernel_size=3, padding='same', activation='relu'),
        layers.Conv2D(512, kernel_size=3, padding='same', activation='relu'),
        layers.MaxPooling2D(pool_size=2, strides=2),

        # Flatten
        layers.Flatten(),

        # Fully Connected Layers
        layers.Dense(4096, activation='relu'),
        layers.Dropout(0.5),
        layers.Dense(4096, activation='relu'),
        layers.Dropout(0.5),
        layers.Dense(num_classes)
    ])

    return model


# ================== Step 3: 工具函数与配置 ==================
def setup_logger(log_dir="logs"):
    os.makedirs(log_dir, exist_ok=True)
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(levelname)s - %(message)s',
        handlers=[
            logging.FileHandler(f"{log_dir}/train_{time.strftime('%Y%m%d_%H%M%S')}.log"),
            logging.StreamHandler()
        ]
    )


# ================== Step 4: 主程序入口 ==================
def main():
    setup_logger()
    strategy = tf.distribute.MirroredStrategy()
    logging.info(f'Number of devices: {strategy.num_replicas_in_sync}')

    with strategy.scope():
        batch_size = 64
        epochs = 50
        num_classes = 10

        # 获取数据集
        train_dataset, val_dataset = get_dataset(batch_size)

        # 构建模型
        model = build_model(num_classes)

        # 编译模型
        optimizer = optimizers.Adam(learning_rate=1e-4)
        loss_fn = losses.SparseCategoricalCrossentropy(from_logits=True)
        metric = metrics.SparseCategoricalAccuracy(name='acc')

        model.compile(optimizer=optimizer,
                      loss=loss_fn,
                      metrics=[metric])

        # 回调函数
        checkpoint_cb = callbacks.ModelCheckpoint(
            "checkpoints/best_vgg16.keras",
            monitor='val_acc',
            save_best_only=True,
            mode='max'
        )
        early_stop_cb = callbacks.EarlyStopping(patience=10, restore_best_weights=True)
        reduce_lr = callbacks.ReduceLROnPlateau(monitor='val_loss', factor=0.2, patience=5)
        tensorboard_cb = callbacks.TensorBoard(log_dir="runs")

        callbacks_list = [checkpoint_cb, early_stop_cb, reduce_lr, tensorboard_cb]

        # 开始训练
        history = model.fit(
            train_dataset,
            validation_data=val_dataset,
            epochs=epochs,
            callbacks=callbacks_list
        )

        logging.info("Training completed.")
        model.save("checkpoints/final_vgg16.keras")


if __name__ == "__main__":
    main()

输出说明

文件夹内容
logs/存放训练日志文件
runs/TensorBoard 日志,用于可视化训练过程
checkpoints/保存最佳模型 best_vgg16.keras 和最终模型 final_vgg16.keras

关键说明

模型结构说明(简化版 VGG16)

层级类型参数说明
Conv2D(64, 3×3)卷积层64 个卷积核,ReLU 激活
MaxPool池化层下采样
Conv2D(128, 3×3)卷积层更深层特征提取
多个卷积层提取高维语义信息
Dense(4096)全连接层高维表示
Dropout正则化防止过拟合
Dense(10)输出层分类输出(logits)

数据预处理

  • 使用了标准的 CIFAR-10 数据集。
  • 图像大小为 (32, 32, 3),已归一化到 [0, 1] 范围。
  • 使用了 tf.data.Dataset 进行高效加载和批处理,并应用了基本的数据增强技术(如随机水平翻转和亮度调整)。

模型编译

  • 使用 Adam 优化器,学习率默认为 1e-4
  • 损失函数为 SparseCategoricalCrossentropy(适用于整数标签)
  • 指标为准确率 SparseCategoricalAccuracy

训练策略

  • 使用了 ModelCheckpoint 保存最佳模型
  • 使用了 ReduceLROnPlateau 动态调整学习率
  • 使用了 EarlyStopping 防止过拟合
  • 使用了 TensorBoard 可视化训练过程

扩展建议

功能描述
增加图像增强添加 RandomFlipRandomRotation 等数据增强操作
替换为原始 VGG16将输入改为 224×224,适配 ImageNet
使用 BatchNormalization提升收敛速度和泛化能力
使用预训练模型如 Keras Applications 中的 ResNet、EfficientNet 等
多 GPU 训练使用 MirroredStrategy 加速训练
自定义损失函数如 Focal Loss、Label Smoothing
模型推理加载模型后进行预测或部署

GoogLeNet / Inception

下面是一个基于 TensorFlow 通用训练模板实现的经典深度学习模型 —— GoogLeNet(Inception v1) 的完整代码示例。我们将使用 CIFAR-10 数据集 来进行训练和验证。由于 GoogLeNet 是为 ImageNet 设计的(输入大小为 224×224),为了适应 CIFAR-10(32×32),我们对网络结构进行了适当的调整。

使用 TensorFlow 通用模板实现 GoogLeNet (Inception v1)

import tensorflow as tf
from tensorflow.keras import layers, models, optimizers, callbacks, losses, metrics
import os
import logging
import time


# ================== Step 1: 数据准备 ==================
def get_dataset(batch_size=64):
    """
    加载并预处理 CIFAR-10 数据集
    """
    (x_train, y_train), (x_val, y_val) = tf.keras.datasets.cifar10.load_data()

    # 归一化 + 类型转换
    x_train = x_train.astype("float32") / 255.0
    x_val = x_val.astype("float32") / 255.0

    train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train))
    val_dataset = tf.data.Dataset.from_tensor_slices((x_val, y_val))

    # 数据增强(可选)
    def preprocess(x, y):
        x = tf.image.random_flip_left_right(x)
        x = tf.image.random_brightness(x, max_delta=0.2)
        return x, tf.squeeze(y)

    train_dataset = train_dataset.map(preprocess).shuffle(10000).batch(batch_size).prefetch(tf.data.AUTOTUNE)
    val_dataset = val_dataset.map(lambda x, y: (x, tf.squeeze(y))).batch(batch_size).prefetch(tf.data.AUTOTUNE)

    return train_dataset, val_dataset


# ================== Step 2: 构建 Inception 模块 ==================
class InceptionModule(layers.Layer):
    def __init__(self, filters_branches, **kwargs):
        super(InceptionModule, self).__init__(**kwargs)
        filter1, filter3_reduce, filter3, filter5_reduce, filter5, pool_proj = filters_branches

        self.branch1 = layers.Conv2D(filter1, kernel_size=1, activation='relu', padding='same')
        self.branch2 = tf.keras.Sequential([
            layers.Conv2D(filter3_reduce, kernel_size=1, activation='relu', padding='same'),
            layers.Conv2D(filter3, kernel_size=3, activation='relu', padding='same')
        ])
        self.branch3 = tf.keras.Sequential([
            layers.Conv2D(filter5_reduce, kernel_size=1, activation='relu', padding='same'),
            layers.Conv2D(filter5, kernel_size=5, activation='relu', padding='same')
        ])
        self.branch4 = tf.keras.Sequential([
            layers.MaxPooling2D(pool_size=3, strides=1, padding='same'),
            layers.Conv2D(pool_proj, kernel_size=1, activation='relu', padding='same')
        ])

    def call(self, inputs, **kwargs):
        branch1 = self.branch1(inputs)
        branch2 = self.branch2(inputs)
        branch3 = self.branch3(inputs)
        branch4 = self.branch4(inputs)
        outputs = layers.concatenate([branch1, branch2, branch3, branch4])
        return outputs


# ================== Step 3: 构建 GoogLeNet 模型 ==================
def build_model(num_classes=10):
    """
    构建简化版 GoogLeNet 模型(适配 CIFAR-10 输入尺寸 32x32x3)
    """
    model = models.Sequential([
        layers.Input(shape=(32, 32, 3)),

        # Initial Convolutional Layer
        layers.Conv2D(64, kernel_size=7, strides=2, padding='same', activation='relu'),
        layers.MaxPooling2D(pool_size=3, strides=2, padding='same'),

        # Inception Module 1
        InceptionModule(filters_branches=[64, 96, 128, 16, 32, 32]),
        InceptionModule(filters_branches=[128, 128, 192, 32, 96, 64]),
        layers.MaxPooling2D(pool_size=3, strides=2, padding='same'),

        # Inception Module 2
        InceptionModule(filters_branches=[192, 96, 208, 16, 48, 64]),
        InceptionModule(filters_branches=[160, 112, 224, 24, 64, 64]),

        # Flatten
        layers.Flatten(),

        # Fully Connected Layers
        layers.Dense(1024, activation='relu'),
        layers.Dropout(0.4),
        layers.Dense(num_classes)
    ])

    return model


# ================== Step 4: 工具函数与配置 ==================
def setup_logger(log_dir="logs"):
    os.makedirs(log_dir, exist_ok=True)
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(levelname)s - %(message)s',
        handlers=[
            logging.FileHandler(f"{log_dir}/train_{time.strftime('%Y%m%d_%H%M%S')}.log"),
            logging.StreamHandler()
        ]
    )


# ================== Step 5: 主程序入口 ==================
def main():
    setup_logger()
    strategy = tf.distribute.MirroredStrategy()
    logging.info(f'Number of devices: {strategy.num_replicas_in_sync}')

    with strategy.scope():
        batch_size = 64
        epochs = 50
        num_classes = 10

        # 获取数据集
        train_dataset, val_dataset = get_dataset(batch_size)

        # 构建模型
        model = build_model(num_classes)

        # 编译模型
        optimizer = optimizers.Adam(learning_rate=1e-4)
        loss_fn = losses.SparseCategoricalCrossentropy(from_logits=True)
        metric = metrics.SparseCategoricalAccuracy(name='acc')

        model.compile(optimizer=optimizer,
                      loss=loss_fn,
                      metrics=[metric])

        # 回调函数
        checkpoint_cb = callbacks.ModelCheckpoint(
            "checkpoints/best_googlenet.keras",
            monitor='val_acc',
            save_best_only=True,
            mode='max'
        )
        early_stop_cb = callbacks.EarlyStopping(patience=10, restore_best_weights=True)
        reduce_lr = callbacks.ReduceLROnPlateau(monitor='val_loss', factor=0.2, patience=5)
        tensorboard_cb = callbacks.TensorBoard(log_dir="runs")

        callbacks_list = [checkpoint_cb, early_stop_cb, reduce_lr, tensorboard_cb]

        # 开始训练
        history = model.fit(
            train_dataset,
            validation_data=val_dataset,
            epochs=epochs,
            callbacks=callbacks_list
        )

        logging.info("Training completed.")
        model.save("checkpoints/final_googlenet.keras")


if __name__ == "__main__":
    main()

输出说明

文件夹内容
logs/存放训练日志文件
runs/TensorBoard 日志,用于可视化训练过程
checkpoints/保存最佳模型 best_googlenet.keras 和最终模型 final_googlenet.keras

关键说明

模型结构说明(简化版 GoogLeNet)

  • Inception 模块:通过并行使用不同大小的卷积核(1×1、3×3、5×5)和最大池化层来提取特征,然后将这些结果拼接起来形成输出。
  • Initial Convolutional Layer:初始卷积层,用于初步特征提取。
  • Fully Connected Layers:全连接层,用于分类任务。

数据预处理

  • 使用了标准的 CIFAR-10 数据集。
  • 图像大小为 (32, 32, 3),已归一化到 [0, 1] 范围。
  • 使用了 tf.data.Dataset 进行高效加载和批处理,并应用了基本的数据增强技术(如随机水平翻转和亮度调整)。

模型编译

  • 使用 Adam 优化器,学习率默认为 1e-4
  • 损失函数为 SparseCategoricalCrossentropy(适用于整数标签)
  • 指标为准确率 SparseCategoricalAccuracy

训练策略

  • 使用了 ModelCheckpoint 保存最佳模型
  • 使用了 ReduceLROnPlateau 动态调整学习率
  • 使用了 EarlyStopping 防止过拟合
  • 使用了 TensorBoard 可视化训练过程

扩展建议

功能描述
增加图像增强添加 RandomFlipRandomRotation 等数据增强操作
替换为原始 GoogLeNet将输入改为 224×224,适配 ImageNet
使用 BatchNormalization提升收敛速度和泛化能力
多 GPU 训练使用 MirroredStrategy 加速训练
自定义损失函数如 Focal Loss、Label Smoothing
模型推理加载模型后进行预测或部署

ResNet

下面是一个基于 TensorFlow 通用训练模板实现的经典深度学习模型 —— ResNet-34 的完整代码示例。我们将使用 CIFAR-10 数据集 来进行训练和验证。由于 ResNet 是为 ImageNet 设计的(输入大小为 224×224),为了适应 CIFAR-10(32×32),我们对网络结构进行了适当的调整。

使用 TensorFlow 通用模板实现 ResNet-34

import tensorflow as tf
from tensorflow.keras import layers, models, optimizers, callbacks, losses, metrics
import os
import logging
import time


# ================== Step 1: 数据准备 ==================
def get_dataset(batch_size=64):
    """
    加载并预处理 CIFAR-10 数据集
    """
    (x_train, y_train), (x_val, y_val) = tf.keras.datasets.cifar10.load_data()

    # 归一化 + 类型转换
    x_train = x_train.astype("float32") / 255.0
    x_val = x_val.astype("float32") / 255.0

    train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train))
    val_dataset = tf.data.Dataset.from_tensor_slices((x_val, y_val))

    # 数据增强(可选)
    def preprocess(x, y):
        x = tf.image.random_flip_left_right(x)
        x = tf.image.random_brightness(x, max_delta=0.2)
        return x, tf.squeeze(y)

    train_dataset = train_dataset.map(preprocess).shuffle(10000).batch(batch_size).prefetch(tf.data.AUTOTUNE)
    val_dataset = val_dataset.map(lambda x, y: (x, tf.squeeze(y))).batch(batch_size).prefetch(tf.data.AUTOTUNE)

    return train_dataset, val_dataset


# ================== Step 2: 构建残差块 ==================
class ResidualBlock(layers.Layer):
    def __init__(self, filters, strides=1, use_projection=False, **kwargs):
        super(ResidualBlock, self).__init__(**kwargs)
        self.use_projection = use_projection

        self.conv1 = layers.Conv2D(filters, kernel_size=3, strides=strides, padding='same')
        self.bn1 = layers.BatchNormalization()
        self.relu1 = layers.ReLU()

        self.conv2 = layers.Conv2D(filters, kernel_size=3, padding='same')
        self.bn2 = layers.BatchNormalization()

        if self.use_projection:
            self.shortcut = layers.Conv2D(filters, kernel_size=1, strides=strides, padding='same')

        self.add = layers.Add()
        self.relu_out = layers.ReLU()

    def call(self, inputs, training=None):
        shortcut = inputs
        if self.use_projection:
            shortcut = self.shortcut(shortcut)

        x = self.conv1(inputs)
        x = self.bn1(x, training=training)
        x = self.relu1(x)

        x = self.conv2(x)
        x = self.bn2(x, training=training)

        x = self.add([x, shortcut])
        return self.relu_out(x)


# ================== Step 3: 构建 ResNet-34 模型 ==================
def build_model(num_classes=10):
    """
    构建简化版 ResNet-34 模型(适配 CIFAR-10 输入尺寸 32x32x3)
    """
    inputs = layers.Input(shape=(32, 32, 3))

    # Initial Convolutional Layer
    x = layers.Conv2D(64, kernel_size=7, strides=2, padding='same')(inputs)
    x = layers.BatchNormalization()(x)
    x = layers.ReLU()(x)
    x = layers.MaxPooling2D(pool_size=3, strides=2, padding='same')(x)

    # Residual Blocks
    def residual_block_sequence(filters, num_blocks, strides_first_block=1):
        x = ResidualBlock(filters=filters, strides=strides_first_block, use_projection=True)(x)
        for _ in range(1, num_blocks):
            x = ResidualBlock(filters=filters)(x)
        return x

    x = residual_block_sequence(64, num_blocks=3, strides_first_block=1)
    x = residual_block_sequence(128, num_blocks=4, strides_first_block=2)
    x = residual_block_sequence(256, num_blocks=6, strides_first_block=2)
    x = residual_block_sequence(512, num_blocks=3, strides_first_block=2)

    # Final Layers
    x = layers.GlobalAveragePooling2D()(x)
    outputs = layers.Dense(num_classes)(x)

    model = models.Model(inputs=inputs, outputs=outputs)
    return model


# ================== Step 4: 工具函数与配置 ==================
def setup_logger(log_dir="logs"):
    os.makedirs(log_dir, exist_ok=True)
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(levelname)s - %(message)s',
        handlers=[
            logging.FileHandler(f"{log_dir}/train_{time.strftime('%Y%m%d_%H%M%S')}.log"),
            logging.StreamHandler()
        ]
    )


# ================== Step 5: 主程序入口 ==================
def main():
    setup_logger()
    strategy = tf.distribute.MirroredStrategy()
    logging.info(f'Number of devices: {strategy.num_replicas_in_sync}')

    with strategy.scope():
        batch_size = 64
        epochs = 50
        num_classes = 10

        # 获取数据集
        train_dataset, val_dataset = get_dataset(batch_size)

        # 构建模型
        model = build_model(num_classes)

        # 编译模型
        optimizer = optimizers.Adam(learning_rate=1e-4)
        loss_fn = losses.SparseCategoricalCrossentropy(from_logits=True)
        metric = metrics.SparseCategoricalAccuracy(name='acc')

        model.compile(optimizer=optimizer,
                      loss=loss_fn,
                      metrics=[metric])

        # 回调函数
        checkpoint_cb = callbacks.ModelCheckpoint(
            "checkpoints/best_resnet34.keras",
            monitor='val_acc',
            save_best_only=True,
            mode='max'
        )
        early_stop_cb = callbacks.EarlyStopping(patience=10, restore_best_weights=True)
        reduce_lr = callbacks.ReduceLROnPlateau(monitor='val_loss', factor=0.2, patience=5)
        tensorboard_cb = callbacks.TensorBoard(log_dir="runs")

        callbacks_list = [checkpoint_cb, early_stop_cb, reduce_lr, tensorboard_cb]

        # 开始训练
        history = model.fit(
            train_dataset,
            validation_data=val_dataset,
            epochs=epochs,
            callbacks=callbacks_list
        )

        logging.info("Training completed.")
        model.save("checkpoints/final_resnet34.keras")


if __name__ == "__main__":
    main()

输出说明

文件夹内容
logs/存放训练日志文件
runs/TensorBoard 日志,用于可视化训练过程
checkpoints/保存最佳模型 best_resnet34.keras 和最终模型 final_resnet34.keras

关键说明

模型结构说明(简化版 ResNet-34)

  • 残差块(Residual Block):通过跳过连接(skip connections)将输入直接加到卷积层输出上,解决了深层网络中的梯度消失问题。
  • Initial Convolutional Layer:初始卷积层,用于初步特征提取。
  • Final Layers:全局平均池化层和全连接层,用于分类任务。

数据预处理

  • 使用了标准的 CIFAR-10 数据集。
  • 图像大小为 (32, 32, 3),已归一化到 [0, 1] 范围。
  • 使用了 tf.data.Dataset 进行高效加载和批处理,并应用了基本的数据增强技术(如随机水平翻转和亮度调整)。

模型编译

  • 使用 Adam 优化器,学习率默认为 1e-4
  • 损失函数为 SparseCategoricalCrossentropy(适用于整数标签)
  • 指标为准确率 SparseCategoricalAccuracy

训练策略

  • 使用了 ModelCheckpoint 保存最佳模型
  • 使用了 ReduceLROnPlateau 动态调整学习率
  • 使用了 EarlyStopping 防止过拟合
  • 使用了 TensorBoard 可视化训练过程

扩展建议

功能描述
增加图像增强添加 RandomFlipRandomRotation 等数据增强操作
替换为原始 ResNet将输入改为 224×224,适配 ImageNet
使用 BatchNormalization提升收敛速度和泛化能力
多 GPU 训练使用 MirroredStrategy 加速训练
自定义损失函数如 Focal Loss、Label Smoothing
模型推理加载模型后进行预测或部署

U-Net

U-Net 是一种常用于医学图像分割的卷积神经网络模型。它特别适合于小数据集训练,通过使用数据增强、加权损失函数和高效的下采样与上采样路径来实现精确的像素级分类。

下面是一个基于 TensorFlow 通用训练模板实现 U-Net 模型,并以一个假设的数据集为例进行训练的完整代码示例。我们将不具体指定数据集,但会提供一个框架,您可以根据自己的数据集(如医学影像分割任务)调整输入部分。

使用 TensorFlow 通用模板实现 U-Net

import tensorflow as tf
from tensorflow.keras import layers, models, optimizers, callbacks, losses
import numpy as np
import os
import logging
import time


# ================== Step 1: 数据准备(这里以随机生成的数据为例)==================
def generate_random_dataset(num_samples=1000, img_shape=(128, 128, 1), mask_shape=(128, 128, 1)):
    """
    随机生成数据集供演示使用。
    """
    images = np.random.rand(num_samples, *img_shape)
    masks = np.random.randint(0, 2, size=(num_samples, *mask_shape))
    return images, masks

def get_dataset(batch_size=32):
    """
    加载并预处理数据集
    """
    images, masks = generate_random_dataset()

    dataset = tf.data.Dataset.from_tensor_slices((images, masks))
    dataset = dataset.shuffle(buffer_size=1024).batch(batch_size).prefetch(tf.data.AUTOTUNE)

    # 分割为训练集和验证集
    train_size = int(0.8 * len(images))
    train_dataset = dataset.take(train_size)
    val_dataset = dataset.skip(train_size)

    return train_dataset, val_dataset


# ================== Step 2: 构建 U-Net 模型 ==================
def conv_block(input_tensor, num_filters):
    x = layers.Conv2D(num_filters, (3, 3), padding='same')(input_tensor)
    x = layers.BatchNormalization()(x)
    x = layers.Activation('relu')(x)
    x = layers.Conv2D(num_filters, (3, 3), padding='same')(x)
    x = layers.BatchNormalization()(x)
    x = layers.Activation('relu')(x)
    return x

def encoder_block(input_tensor, num_filters):
    x = conv_block(input_tensor, num_filters)
    p = layers.MaxPooling2D((2, 2))(x)
    return x, p

def decoder_block(input_tensor, skip_features, num_filters):
    x = layers.Conv2DTranspose(num_filters, (2, 2), strides=2, padding='same')(input_tensor)
    x = layers.concatenate([x, skip_features])
    x = conv_block(x, num_filters)
    return x

def build_unet(img_shape):
    inputs = layers.Input(shape=img_shape)

    s1, p1 = encoder_block(inputs, 64)
    s2, p2 = encoder_block(p1, 128)
    s3, p3 = encoder_block(p2, 256)
    s4, p4 = encoder_block(p3, 512)

    b = conv_block(p4, 1024)

    d1 = decoder_block(b, s4, 512)
    d2 = decoder_block(d1, s3, 256)
    d3 = decoder_block(d2, s2, 128)
    d4 = decoder_block(d3, s1, 64)

    outputs = layers.Conv2D(1, (1, 1), activation='sigmoid')(d4)

    model = models.Model(inputs=[inputs], outputs=[outputs])
    return model


# ================== Step 3: 工具函数与配置 ==================
def setup_logger(log_dir="logs"):
    os.makedirs(log_dir, exist_ok=True)
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(levelname)s - %(message)s',
        handlers=[
            logging.FileHandler(f"{log_dir}/train_{time.strftime('%Y%m%d_%H%M%S')}.log"),
            logging.StreamHandler()
        ]
    )


# ================== Step 4: 主程序入口 ==================
def main():
    setup_logger()
    strategy = tf.distribute.MirroredStrategy()
    logging.info(f'Number of devices: {strategy.num_replicas_in_sync}')

    with strategy.scope():
        batch_size = 32
        epochs = 50
        img_shape = (128, 128, 1)

        # 获取数据集
        train_dataset, val_dataset = get_dataset(batch_size)

        # 构建模型
        model = build_unet(img_shape)

        # 编译模型
        optimizer = optimizers.Adam(learning_rate=1e-4)
        loss_fn = losses.BinaryCrossentropy(from_logits=False)
        metric = 'accuracy'

        model.compile(optimizer=optimizer,
                      loss=loss_fn,
                      metrics=[metric])

        # 回调函数
        checkpoint_cb = callbacks.ModelCheckpoint(
            "checkpoints/best_unet.keras",
            monitor='val_accuracy',
            save_best_only=True,
            mode='max'
        )
        early_stop_cb = callbacks.EarlyStopping(patience=10, restore_best_weights=True)
        reduce_lr = callbacks.ReduceLROnPlateau(monitor='val_loss', factor=0.2, patience=5)
        tensorboard_cb = callbacks.TensorBoard(log_dir="runs")

        callbacks_list = [checkpoint_cb, early_stop_cb, reduce_lr, tensorboard_cb]

        # 开始训练
        history = model.fit(
            train_dataset,
            validation_data=val_dataset,
            epochs=epochs,
            callbacks=callbacks_list
        )

        logging.info("Training completed.")
        model.save("checkpoints/final_unet.keras")


if __name__ == "__main__":
    main()

输出说明

文件夹内容
logs/存放训练日志文件
runs/TensorBoard 日志,用于可视化训练过程
checkpoints/保存最佳模型 best_unet.keras 和最终模型 final_unet.keras

关键说明

模型结构说明(U-Net)

  • 编码器(Encoder):通过一系列的卷积和最大池化操作提取特征图。每个块包括两个卷积层,接着是批量归一化和 ReLU 激活函数。
  • 解码器(Decoder):通过转置卷积操作逐步恢复空间分辨率。每个块接收来自编码器对应阶段的特征映射作为跳过连接,以帮助保留位置信息。
  • 输出层:应用一个卷积层将通道数减少到所需的类别数量(本例中为 1),并使用 Sigmoid 激活函数进行二分类。

数据预处理

  • 示例中使用了随机生成的数据集。在实际应用中,您需要替换这部分代码以加载和预处理您的特定数据集,例如医学影像数据。

模型编译

  • 使用 Adam 优化器,学习率默认为 1e-4
  • 损失函数为 BinaryCrossentropy(适用于二分类问题)
  • 指标为准确率

训练策略

  • 使用了 ModelCheckpoint 保存最佳模型
  • 使用了 ReduceLROnPlateau 动态调整学习率
  • 使用了 EarlyStopping 防止过拟合
  • 使用了 TensorBoard 可视化训练过程

扩展建议

功能描述
增加图像增强添加 RandomFlipRandomRotation 等数据增强操作
替换为真实数据集如医学影像数据集等
多 GPU 训练使用 MirroredStrategy 加速训练
自定义损失函数如 Dice Loss 或 Jaccard Index 等更适合分割任务的损失函数
模型推理加载模型后进行预测或部署

GANs (DCGAN)

下面是一个基于 TensorFlow 通用训练模板实现的经典 GAN(生成对抗网络) 的完整代码示例。我们将使用 DCGAN(深度卷积 GAN) 架构,在 MNIST 手写数字数据集 上进行训练,以演示如何从随机噪声生成逼真的图像。

使用 TensorFlow 通用模板实现 DCGAN

import tensorflow as tf
from tensorflow.keras import layers, models, optimizers, callbacks
from tensorflow.keras.datasets import mnist
import numpy as np
import os
import logging
import time
import matplotlib.pyplot as plt


# ================== Step 1: 数据准备 ==================
def get_dataset(batch_size=128):
    """
    加载并预处理 MNIST 数据集
    """
    (train_images, _), _ = mnist.load_data()
    train_images = train_images.reshape(-1, 28, 28, 1).astype("float32") / 127.5 - 1.0  # [-1, 1] 归一化

    dataset = tf.data.Dataset.from_tensor_slices(train_images)
    dataset = dataset.shuffle(buffer_size=60000).batch(batch_size).prefetch(tf.data.AUTOTUNE)

    return dataset


# ================== Step 2: 构建生成器和判别器 ==================

def build_generator(latent_dim):
    """
    生成器:将随机噪声向量映射为 28x28x1 图像
    """
    model = models.Sequential([
        layers.Dense(7 * 7 * 256, use_bias=False, input_shape=(latent_dim,)),
        layers.BatchNormalization(),
        layers.ReLU(),

        layers.Reshape((7, 7, 256)),

        layers.Conv2DTranspose(128, (5, 5), strides=(1, 1), padding='same', use_bias=False),
        layers.BatchNormalization(),
        layers.ReLU(),

        layers.Conv2DTranspose(64, (5, 5), strides=(2, 2), padding='same', use_bias=False),
        layers.BatchNormalization(),
        layers.ReLU(),

        layers.Conv2DTranspose(1, (5, 5), strides=(2, 2), padding='same', use_bias=False, activation='tanh')
    ])

    return model


def build_discriminator():
    """
    判别器:判断图像是真实的还是生成的
    """
    model = models.Sequential([
        layers.Input(shape=(28, 28, 1)),

        layers.Conv2D(64, (5, 5), strides=(2, 2), padding='same'),
        layers.LeakyReLU(alpha=0.2),

        layers.Conv2D(128, (5, 5), strides=(2, 2), padding='same'),
        layers.LeakyReLU(alpha=0.2),

        layers.Flatten(),
        layers.Dense(1)  # 输出一个实数分数(不加 sigmoid)
    ])

    return model


# ================== Step 3: 定义 GAN 模型 ==================
class DCGAN(models.Model):
    def __init__(self, generator, discriminator, latent_dim, **kwargs):
        super(DCGAN, self).__init__(**kwargs)
        self.generator = generator
        self.discriminator = discriminator
        self.latent_dim = latent_dim

    def compile(self, g_optimizer, d_optimizer, loss_fn):
        super(DCGAN, self).compile()
        self.g_optimizer = g_optimizer
        self.d_optimizer = d_optimizer
        self.loss_fn = loss_fn

    def train_step(self, real_images):

        batch_size = tf.shape(real_images)[0]

        # ----------------------------
        # 训练判别器
        # ----------------------------
        noise = tf.random.normal(shape=(batch_size, self.latent_dim))
        generated_images = self.generator(noise)

        with tf.GradientTape() as tape:
            real_output = self.discriminator(real_images, training=True)
            fake_output = self.discriminator(generated_images, training=True)

            real_loss = self.loss_fn(tf.ones_like(real_output), real_output)
            fake_loss = self.loss_fn(tf.zeros_like(fake_output), fake_output)
            d_loss = real_loss + fake_loss

        grads = tape.gradient(d_loss, self.discriminator.trainable_weights)
        self.d_optimizer.apply_gradients(zip(grads, self.discriminator.trainable_weights))

        # ----------------------------
        # 训练生成器
        # ----------------------------
        noise = tf.random.normal(shape=(batch_size, self.latent_dim))

        with tf.GradientTape() as tape:
            generated_images = self.generator(noise, training=True)
            fake_output = self.discriminator(generated_images, training=False)

            g_loss = self.loss_fn(tf.ones_like(fake_output), fake_output)

        grads = tape.gradient(g_loss, self.generator.trainable_weights)
        self.g_optimizer.apply_gradients(zip(grads, self.generator.trainable_weights))

        return {"d_loss": d_loss, "g_loss": g_loss}


# ================== Step 4: 工具函数与配置 ==================
def setup_logger(log_dir="logs"):
    os.makedirs(log_dir, exist_ok=True)
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(levelname)s - %(message)s',
        handlers=[
            logging.FileHandler(f"{log_dir}/train_{time.strftime('%Y%m%d_%H%M%S')}.log"),
            logging.StreamHandler()
        ]
    )

def plot_generated_images(epoch, generator, test_input, save_path="images"):
    predictions = generator(test_input, training=False)
    fig, axes = plt.subplots(4, 4, figsize=(8, 8))
    for i, ax in enumerate(axes.flat):
        ax.imshow(predictions[i, :, :, 0] * 0.5 + 0.5, cmap='gray')
        ax.axis("off")
    plt.tight_layout()
    os.makedirs(save_path, exist_ok=True)
    plt.savefig(f"{save_path}/image_at_epoch_{epoch:04d}.png")
    plt.close()


# ================== Step 5: 主程序入口 ==================
def main():
    setup_logger()
    strategy = tf.distribute.MirroredStrategy()
    logging.info(f'Number of devices: {strategy.num_replicas_in_sync}')

    with strategy.scope():
        batch_size = 128
        epochs = 100
        latent_dim = 100

        # 获取数据集
        dataset = get_dataset(batch_size)

        # 构建模型
        generator = build_generator(latent_dim)
        discriminator = build_discriminator()

        # 编译模型
        dcgan = DCGAN(generator, discriminator, latent_dim)
        dcgan.compile(
            g_optimizer=optimizers.Adam(learning_rate=2e-4, beta_1=0.5),
            d_optimizer=optimizers.Adam(learning_rate=2e-4, beta_1=0.5),
            loss_fn=losses.Hinge()  # 可选 Hinge Loss 或 BinaryCrossentropy
        )

        # 回调函数
        checkpoint_cb = callbacks.ModelCheckpoint(
            "checkpoints/generator.keras",
            monitor='g_loss',
            save_best_only=True,
            mode='min'
        )
        tensorboard_cb = callbacks.TensorBoard(log_dir="runs")

        fixed_noise = tf.random.normal([16, latent_dim])  # 固定噪声用于可视化生成效果

        # 自定义回调:定期保存生成图像
        class ImageSaveCallback(callbacks.Callback):
            def on_epoch_end(self, epoch, logs=None):
                if (epoch + 1) % 5 == 0:
                    plot_generated_images(epoch + 1, self.model.generator, fixed_noise)

        image_save_cb = ImageSaveCallback()

        callbacks_list = [checkpoint_cb, tensorboard_cb, image_save_cb]

        # 开始训练
        history = dcgan.fit(
            dataset,
            epochs=epochs,
            callbacks=callbacks_list
        )

        logging.info("Training completed.")
        dcgan.generator.save("checkpoints/final_generator.keras")
        dcgan.discriminator.save("checkpoints/final_discriminator.keras")


if __name__ == "__main__":
    main()

输出说明

文件夹内容
logs/存放训练日志文件
runs/TensorBoard 日志,用于可视化训练过程
images/每隔一定轮次保存的生成图像
checkpoints/保存最佳生成器、最终生成器和判别器

关键说明

模型结构说明(DCGAN)

部分类型参数说明
GeneratorConv2DTranspose将噪声向量映射为真实图像
DiscriminatorConv2D判断图像是真实还是生成的
Loss FunctionHinge Loss更适合稳定训练 GAN
OptimizerAdam(2e-4, beta=0.5)常用于 GAN 的优化器参数

数据预处理

  • 使用了标准的 MNIST 数据集
  • 图像大小为 (28, 28, 1),归一化到 [-1, 1]
  • 使用 tf.data.Dataset 进行高效加载和批处理

模型编译与训练

  • 使用自定义 train_step() 来实现交替训练策略
  • 判别器使用真实标签 [1] 和生成标签 [0] 来计算损失
  • 生成器的目标是让判别器输出接近 [1]

扩展建议

功能描述
替换为其他架构如 CycleGAN、StyleGAN、Wasserstein-GAN 等
多 GPU 训练使用 MirroredStrategy 加速训练
自定义损失函数如 Wasserstein loss、LS-GAN loss
图像转译任务如 pix2pix、CycleGAN 等条件 GAN
模型推理用固定噪声生成图像或图像修复等

VGG

下面是一个基于 TensorFlow 通用训练模板实现的经典深度学习模型 —— VGG16 的完整代码示例。我们将使用 CIFAR-10 数据集 来进行训练和验证。由于 VGG16 是为 ImageNet 设计的(输入大小为 224×224),为了适应 CIFAR-10(32×32),我们对网络结构进行了适当的调整。

使用 TensorFlow 通用模板实现 VGG16

import tensorflow as tf
from tensorflow.keras import layers, models, optimizers, callbacks, losses, metrics
import os
import logging
import time


# ================== Step 1: 数据准备 ==================
def get_dataset(batch_size=64):
    """
    加载并预处理 CIFAR-10 数据集
    """
    (x_train, y_train), (x_val, y_val) = tf.keras.datasets.cifar10.load_data()

    # 归一化 + 类型转换
    x_train = x_train.astype("float32") / 255.0
    x_val = x_val.astype("float32") / 255.0

    train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train))
    val_dataset = tf.data.Dataset.from_tensor_slices((x_val, y_val))

    # 数据增强(可选)
    def preprocess(x, y):
        x = tf.image.random_flip_left_right(x)
        x = tf.image.random_brightness(x, max_delta=0.2)
        return x, tf.squeeze(y)

    train_dataset = train_dataset.map(preprocess).shuffle(10000).batch(batch_size).prefetch(tf.data.AUTOTUNE)
    val_dataset = val_dataset.map(lambda x, y: (x, tf.squeeze(y))).batch(batch_size).prefetch(tf.data.AUTOTUNE)

    return train_dataset, val_dataset


# ================== Step 2: 构建 VGG16 模型 ==================
def build_model(num_classes=10):
    """
    构建简化版 VGG16 模型(适配 CIFAR-10 输入尺寸 32x32x3)
    """
    model = models.Sequential([
        layers.Input(shape=(32, 32, 3)),

        # Block 1
        layers.Conv2D(64, kernel_size=3, padding='same', activation='relu'),
        layers.Conv2D(64, kernel_size=3, padding='same', activation='relu'),
        layers.MaxPooling2D(pool_size=2, strides=2),

        # Block 2
        layers.Conv2D(128, kernel_size=3, padding='same', activation='relu'),
        layers.Conv2D(128, kernel_size=3, padding='same', activation='relu'),
        layers.MaxPooling2D(pool_size=2, strides=2),

        # Block 3
        layers.Conv2D(256, kernel_size=3, padding='same', activation='relu'),
        layers.Conv2D(256, kernel_size=3, padding='same', activation='relu'),
        layers.Conv2D(256, kernel_size=3, padding='same', activation='relu'),
        layers.MaxPooling2D(pool_size=2, strides=2),

        # Block 4
        layers.Conv2D(512, kernel_size=3, padding='same', activation='relu'),
        layers.Conv2D(512, kernel_size=3, padding='same', activation='relu'),
        layers.Conv2D(512, kernel_size=3, padding='same', activation='relu'),
        layers.MaxPooling2D(pool_size=2, strides=2),

        # Block 5
        layers.Conv2D(512, kernel_size=3, padding='same', activation='relu'),
        layers.Conv2D(512, kernel_size=3, padding='same', activation='relu'),
        layers.Conv2D(512, kernel_size=3, padding='same', activation='relu'),
        layers.MaxPooling2D(pool_size=2, strides=2),

        # Flatten
        layers.Flatten(),

        # Fully Connected Layers
        layers.Dense(4096, activation='relu'),
        layers.Dropout(0.5),
        layers.Dense(4096, activation='relu'),
        layers.Dropout(0.5),
        layers.Dense(num_classes)
    ])

    return model


# ================== Step 3: 工具函数与配置 ==================
def setup_logger(log_dir="logs"):
    os.makedirs(log_dir, exist_ok=True)
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(levelname)s - %(message)s',
        handlers=[
            logging.FileHandler(f"{log_dir}/train_{time.strftime('%Y%m%d_%H%M%S')}.log"),
            logging.StreamHandler()
        ]
    )


# ================== Step 4: 主程序入口 ==================
def main():
    setup_logger()
    strategy = tf.distribute.MirroredStrategy()
    logging.info(f'Number of devices: {strategy.num_replicas_in_sync}')

    with strategy.scope():
        batch_size = 64
        epochs = 50
        num_classes = 10

        # 获取数据集
        train_dataset, val_dataset = get_dataset(batch_size)

        # 构建模型
        model = build_model(num_classes)

        # 编译模型
        optimizer = optimizers.Adam(learning_rate=1e-4)
        loss_fn = losses.SparseCategoricalCrossentropy(from_logits=True)
        metric = metrics.SparseCategoricalAccuracy(name='acc')

        model.compile(optimizer=optimizer,
                      loss=loss_fn,
                      metrics=[metric])

        # 回调函数
        checkpoint_cb = callbacks.ModelCheckpoint(
            "checkpoints/best_vgg16.keras",
            monitor='val_acc',
            save_best_only=True,
            mode='max'
        )
        early_stop_cb = callbacks.EarlyStopping(patience=10, restore_best_weights=True)
        reduce_lr = callbacks.ReduceLROnPlateau(monitor='val_loss', factor=0.2, patience=5)
        tensorboard_cb = callbacks.TensorBoard(log_dir="runs")

        callbacks_list = [checkpoint_cb, early_stop_cb, reduce_lr, tensorboard_cb]

        # 开始训练
        history = model.fit(
            train_dataset,
            validation_data=val_dataset,
            epochs=epochs,
            callbacks=callbacks_list
        )

        logging.info("Training completed.")
        model.save("checkpoints/final_vgg16.keras")


if __name__ == "__main__":
    main()

输出说明

文件夹内容
logs/存放训练日志文件
runs/TensorBoard 日志,用于可视化训练过程
checkpoints/保存最佳模型 best_vgg16.keras 和最终模型 final_vgg16.keras

关键说明

模型结构说明(简化版 VGG16)

层级类型参数说明
Conv2D(64, 3×3)卷积层64 个卷积核,ReLU 激活
MaxPool池化层下采样
Conv2D(128, 3×3)卷积层更深层特征提取
多个卷积层提取高维语义信息
Dense(4096)全连接层高维表示
Dropout正则化防止过拟合
Dense(10)输出层分类输出(logits)

数据预处理

  • 使用了标准的 CIFAR-10 数据集。
  • 图像大小为 (32, 32, 3),已归一化到 [0, 1] 范围。
  • 使用了 tf.data.Dataset 进行高效加载和批处理,并应用了基本的数据增强技术(如随机水平翻转和亮度调整)。

模型编译

  • 使用 Adam 优化器,学习率默认为 1e-4
  • 损失函数为 SparseCategoricalCrossentropy(适用于整数标签)
  • 指标为准确率 SparseCategoricalAccuracy

训练策略

  • 使用了 ModelCheckpoint 保存最佳模型
  • 使用了 ReduceLROnPlateau 动态调整学习率
  • 使用了 EarlyStopping 防止过拟合
  • 使用了 TensorBoard 可视化训练过程

扩展建议

功能描述
增加图像增强添加 RandomFlipRandomRotation 等数据增强操作
替换为原始 VGG16将输入改为 224×224,适配 ImageNet
使用 BatchNormalization提升收敛速度和泛化能力
使用预训练模型如 Keras Applications 中的 ResNet、EfficientNet 等
多 GPU 训练使用 MirroredStrategy 加速训练
自定义损失函数如 Focal Loss、Label Smoothing
模型推理加载模型后进行预测或部署

GoogLeNet / Inception

下面是一个基于 TensorFlow 通用训练模板实现的经典深度学习模型 —— GoogLeNet(Inception v1) 的完整代码示例。我们将使用 CIFAR-10 数据集 来进行训练和验证。由于 GoogLeNet 是为 ImageNet 设计的(输入大小为 224×224),为了适应 CIFAR-10(32×32),我们对网络结构进行了适当的调整。

使用 TensorFlow 通用模板实现 GoogLeNet (Inception v1)

import tensorflow as tf
from tensorflow.keras import layers, models, optimizers, callbacks, losses, metrics
import os
import logging
import time


# ================== Step 1: 数据准备 ==================
def get_dataset(batch_size=64):
    """
    加载并预处理 CIFAR-10 数据集
    """
    (x_train, y_train), (x_val, y_val) = tf.keras.datasets.cifar10.load_data()

    # 归一化 + 类型转换
    x_train = x_train.astype("float32") / 255.0
    x_val = x_val.astype("float32") / 255.0

    train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train))
    val_dataset = tf.data.Dataset.from_tensor_slices((x_val, y_val))

    # 数据增强(可选)
    def preprocess(x, y):
        x = tf.image.random_flip_left_right(x)
        x = tf.image.random_brightness(x, max_delta=0.2)
        return x, tf.squeeze(y)

    train_dataset = train_dataset.map(preprocess).shuffle(10000).batch(batch_size).prefetch(tf.data.AUTOTUNE)
    val_dataset = val_dataset.map(lambda x, y: (x, tf.squeeze(y))).batch(batch_size).prefetch(tf.data.AUTOTUNE)

    return train_dataset, val_dataset


# ================== Step 2: 构建 Inception 模块 ==================
class InceptionModule(layers.Layer):
    def __init__(self, filters_branches, **kwargs):
        super(InceptionModule, self).__init__(**kwargs)
        filter1, filter3_reduce, filter3, filter5_reduce, filter5, pool_proj = filters_branches

        self.branch1 = layers.Conv2D(filter1, kernel_size=1, activation='relu', padding='same')
        self.branch2 = tf.keras.Sequential([
            layers.Conv2D(filter3_reduce, kernel_size=1, activation='relu', padding='same'),
            layers.Conv2D(filter3, kernel_size=3, activation='relu', padding='same')
        ])
        self.branch3 = tf.keras.Sequential([
            layers.Conv2D(filter5_reduce, kernel_size=1, activation='relu', padding='same'),
            layers.Conv2D(filter5, kernel_size=5, activation='relu', padding='same')
        ])
        self.branch4 = tf.keras.Sequential([
            layers.MaxPooling2D(pool_size=3, strides=1, padding='same'),
            layers.Conv2D(pool_proj, kernel_size=1, activation='relu', padding='same')
        ])

    def call(self, inputs, **kwargs):
        branch1 = self.branch1(inputs)
        branch2 = self.branch2(inputs)
        branch3 = self.branch3(inputs)
        branch4 = self.branch4(inputs)
        outputs = layers.concatenate([branch1, branch2, branch3, branch4])
        return outputs


# ================== Step 3: 构建 GoogLeNet 模型 ==================
def build_model(num_classes=10):
    """
    构建简化版 GoogLeNet 模型(适配 CIFAR-10 输入尺寸 32x32x3)
    """
    model = models.Sequential([
        layers.Input(shape=(32, 32, 3)),

        # Initial Convolutional Layer
        layers.Conv2D(64, kernel_size=7, strides=2, padding='same', activation='relu'),
        layers.MaxPooling2D(pool_size=3, strides=2, padding='same'),

        # Inception Module 1
        InceptionModule(filters_branches=[64, 96, 128, 16, 32, 32]),
        InceptionModule(filters_branches=[128, 128, 192, 32, 96, 64]),
        layers.MaxPooling2D(pool_size=3, strides=2, padding='same'),

        # Inception Module 2
        InceptionModule(filters_branches=[192, 96, 208, 16, 48, 64]),
        InceptionModule(filters_branches=[160, 112, 224, 24, 64, 64]),

        # Flatten
        layers.Flatten(),

        # Fully Connected Layers
        layers.Dense(1024, activation='relu'),
        layers.Dropout(0.4),
        layers.Dense(num_classes)
    ])

    return model


# ================== Step 4: 工具函数与配置 ==================
def setup_logger(log_dir="logs"):
    os.makedirs(log_dir, exist_ok=True)
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(levelname)s - %(message)s',
        handlers=[
            logging.FileHandler(f"{log_dir}/train_{time.strftime('%Y%m%d_%H%M%S')}.log"),
            logging.StreamHandler()
        ]
    )


# ================== Step 5: 主程序入口 ==================
def main():
    setup_logger()
    strategy = tf.distribute.MirroredStrategy()
    logging.info(f'Number of devices: {strategy.num_replicas_in_sync}')

    with strategy.scope():
        batch_size = 64
        epochs = 50
        num_classes = 10

        # 获取数据集
        train_dataset, val_dataset = get_dataset(batch_size)

        # 构建模型
        model = build_model(num_classes)

        # 编译模型
        optimizer = optimizers.Adam(learning_rate=1e-4)
        loss_fn = losses.SparseCategoricalCrossentropy(from_logits=True)
        metric = metrics.SparseCategoricalAccuracy(name='acc')

        model.compile(optimizer=optimizer,
                      loss=loss_fn,
                      metrics=[metric])

        # 回调函数
        checkpoint_cb = callbacks.ModelCheckpoint(
            "checkpoints/best_googlenet.keras",
            monitor='val_acc',
            save_best_only=True,
            mode='max'
        )
        early_stop_cb = callbacks.EarlyStopping(patience=10, restore_best_weights=True)
        reduce_lr = callbacks.ReduceLROnPlateau(monitor='val_loss', factor=0.2, patience=5)
        tensorboard_cb = callbacks.TensorBoard(log_dir="runs")

        callbacks_list = [checkpoint_cb, early_stop_cb, reduce_lr, tensorboard_cb]

        # 开始训练
        history = model.fit(
            train_dataset,
            validation_data=val_dataset,
            epochs=epochs,
            callbacks=callbacks_list
        )

        logging.info("Training completed.")
        model.save("checkpoints/final_googlenet.keras")


if __name__ == "__main__":
    main()

输出说明

文件夹内容
logs/存放训练日志文件
runs/TensorBoard 日志,用于可视化训练过程
checkpoints/保存最佳模型 best_googlenet.keras 和最终模型 final_googlenet.keras

关键说明

模型结构说明(简化版 GoogLeNet)

  • Inception 模块:通过并行使用不同大小的卷积核(1×1、3×3、5×5)和最大池化层来提取特征,然后将这些结果拼接起来形成输出。
  • Initial Convolutional Layer:初始卷积层,用于初步特征提取。
  • Fully Connected Layers:全连接层,用于分类任务。

数据预处理

  • 使用了标准的 CIFAR-10 数据集。
  • 图像大小为 (32, 32, 3),已归一化到 [0, 1] 范围。
  • 使用了 tf.data.Dataset 进行高效加载和批处理,并应用了基本的数据增强技术(如随机水平翻转和亮度调整)。

模型编译

  • 使用 Adam 优化器,学习率默认为 1e-4
  • 损失函数为 SparseCategoricalCrossentropy(适用于整数标签)
  • 指标为准确率 SparseCategoricalAccuracy

训练策略

  • 使用了 ModelCheckpoint 保存最佳模型
  • 使用了 ReduceLROnPlateau 动态调整学习率
  • 使用了 EarlyStopping 防止过拟合
  • 使用了 TensorBoard 可视化训练过程

扩展建议

功能描述
增加图像增强添加 RandomFlipRandomRotation 等数据增强操作
替换为原始 GoogLeNet将输入改为 224×224,适配 ImageNet
使用 BatchNormalization提升收敛速度和泛化能力
多 GPU 训练使用 MirroredStrategy 加速训练
自定义损失函数如 Focal Loss、Label Smoothing
模型推理加载模型后进行预测或部署

ResNet

下面是一个基于 TensorFlow 通用训练模板实现的经典深度学习模型 —— ResNet-34 的完整代码示例。我们将使用 CIFAR-10 数据集 来进行训练和验证。由于 ResNet 是为 ImageNet 设计的(输入大小为 224×224),为了适应 CIFAR-10(32×32),我们对网络结构进行了适当的调整。

使用 TensorFlow 通用模板实现 ResNet-34

import tensorflow as tf
from tensorflow.keras import layers, models, optimizers, callbacks, losses, metrics
import os
import logging
import time


# ================== Step 1: 数据准备 ==================
def get_dataset(batch_size=64):
    """
    加载并预处理 CIFAR-10 数据集
    """
    (x_train, y_train), (x_val, y_val) = tf.keras.datasets.cifar10.load_data()

    # 归一化 + 类型转换
    x_train = x_train.astype("float32") / 255.0
    x_val = x_val.astype("float32") / 255.0

    train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train))
    val_dataset = tf.data.Dataset.from_tensor_slices((x_val, y_val))

    # 数据增强(可选)
    def preprocess(x, y):
        x = tf.image.random_flip_left_right(x)
        x = tf.image.random_brightness(x, max_delta=0.2)
        return x, tf.squeeze(y)

    train_dataset = train_dataset.map(preprocess).shuffle(10000).batch(batch_size).prefetch(tf.data.AUTOTUNE)
    val_dataset = val_dataset.map(lambda x, y: (x, tf.squeeze(y))).batch(batch_size).prefetch(tf.data.AUTOTUNE)

    return train_dataset, val_dataset


# ================== Step 2: 构建残差块 ==================
class ResidualBlock(layers.Layer):
    def __init__(self, filters, strides=1, use_projection=False, **kwargs):
        super(ResidualBlock, self).__init__(**kwargs)
        self.use_projection = use_projection

        self.conv1 = layers.Conv2D(filters, kernel_size=3, strides=strides, padding='same')
        self.bn1 = layers.BatchNormalization()
        self.relu1 = layers.ReLU()

        self.conv2 = layers.Conv2D(filters, kernel_size=3, padding='same')
        self.bn2 = layers.BatchNormalization()

        if self.use_projection:
            self.shortcut = layers.Conv2D(filters, kernel_size=1, strides=strides, padding='same')

        self.add = layers.Add()
        self.relu_out = layers.ReLU()

    def call(self, inputs, training=None):
        shortcut = inputs
        if self.use_projection:
            shortcut = self.shortcut(shortcut)

        x = self.conv1(inputs)
        x = self.bn1(x, training=training)
        x = self.relu1(x)

        x = self.conv2(x)
        x = self.bn2(x, training=training)

        x = self.add([x, shortcut])
        return self.relu_out(x)


# ================== Step 3: 构建 ResNet-34 模型 ==================
def build_model(num_classes=10):
    """
    构建简化版 ResNet-34 模型(适配 CIFAR-10 输入尺寸 32x32x3)
    """
    inputs = layers.Input(shape=(32, 32, 3))

    # Initial Convolutional Layer
    x = layers.Conv2D(64, kernel_size=7, strides=2, padding='same')(inputs)
    x = layers.BatchNormalization()(x)
    x = layers.ReLU()(x)
    x = layers.MaxPooling2D(pool_size=3, strides=2, padding='same')(x)

    # Residual Blocks
    def residual_block_sequence(filters, num_blocks, strides_first_block=1):
        x = ResidualBlock(filters=filters, strides=strides_first_block, use_projection=True)(x)
        for _ in range(1, num_blocks):
            x = ResidualBlock(filters=filters)(x)
        return x

    x = residual_block_sequence(64, num_blocks=3, strides_first_block=1)
    x = residual_block_sequence(128, num_blocks=4, strides_first_block=2)
    x = residual_block_sequence(256, num_blocks=6, strides_first_block=2)
    x = residual_block_sequence(512, num_blocks=3, strides_first_block=2)

    # Final Layers
    x = layers.GlobalAveragePooling2D()(x)
    outputs = layers.Dense(num_classes)(x)

    model = models.Model(inputs=inputs, outputs=outputs)
    return model


# ================== Step 4: 工具函数与配置 ==================
def setup_logger(log_dir="logs"):
    os.makedirs(log_dir, exist_ok=True)
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(levelname)s - %(message)s',
        handlers=[
            logging.FileHandler(f"{log_dir}/train_{time.strftime('%Y%m%d_%H%M%S')}.log"),
            logging.StreamHandler()
        ]
    )


# ================== Step 5: 主程序入口 ==================
def main():
    setup_logger()
    strategy = tf.distribute.MirroredStrategy()
    logging.info(f'Number of devices: {strategy.num_replicas_in_sync}')

    with strategy.scope():
        batch_size = 64
        epochs = 50
        num_classes = 10

        # 获取数据集
        train_dataset, val_dataset = get_dataset(batch_size)

        # 构建模型
        model = build_model(num_classes)

        # 编译模型
        optimizer = optimizers.Adam(learning_rate=1e-4)
        loss_fn = losses.SparseCategoricalCrossentropy(from_logits=True)
        metric = metrics.SparseCategoricalAccuracy(name='acc')

        model.compile(optimizer=optimizer,
                      loss=loss_fn,
                      metrics=[metric])

        # 回调函数
        checkpoint_cb = callbacks.ModelCheckpoint(
            "checkpoints/best_resnet34.keras",
            monitor='val_acc',
            save_best_only=True,
            mode='max'
        )
        early_stop_cb = callbacks.EarlyStopping(patience=10, restore_best_weights=True)
        reduce_lr = callbacks.ReduceLROnPlateau(monitor='val_loss', factor=0.2, patience=5)
        tensorboard_cb = callbacks.TensorBoard(log_dir="runs")

        callbacks_list = [checkpoint_cb, early_stop_cb, reduce_lr, tensorboard_cb]

        # 开始训练
        history = model.fit(
            train_dataset,
            validation_data=val_dataset,
            epochs=epochs,
            callbacks=callbacks_list
        )

        logging.info("Training completed.")
        model.save("checkpoints/final_resnet34.keras")


if __name__ == "__main__":
    main()

输出说明

文件夹内容
logs/存放训练日志文件
runs/TensorBoard 日志,用于可视化训练过程
checkpoints/保存最佳模型 best_resnet34.keras 和最终模型 final_resnet34.keras

关键说明

模型结构说明(简化版 ResNet-34)

  • 残差块(Residual Block):通过跳过连接(skip connections)将输入直接加到卷积层输出上,解决了深层网络中的梯度消失问题。
  • Initial Convolutional Layer:初始卷积层,用于初步特征提取。
  • Final Layers:全局平均池化层和全连接层,用于分类任务。

数据预处理

  • 使用了标准的 CIFAR-10 数据集。
  • 图像大小为 (32, 32, 3),已归一化到 [0, 1] 范围。
  • 使用了 tf.data.Dataset 进行高效加载和批处理,并应用了基本的数据增强技术(如随机水平翻转和亮度调整)。

模型编译

  • 使用 Adam 优化器,学习率默认为 1e-4
  • 损失函数为 SparseCategoricalCrossentropy(适用于整数标签)
  • 指标为准确率 SparseCategoricalAccuracy

训练策略

  • 使用了 ModelCheckpoint 保存最佳模型
  • 使用了 ReduceLROnPlateau 动态调整学习率
  • 使用了 EarlyStopping 防止过拟合
  • 使用了 TensorBoard 可视化训练过程

扩展建议

功能描述
增加图像增强添加 RandomFlipRandomRotation 等数据增强操作
替换为原始 ResNet将输入改为 224×224,适配 ImageNet
使用 BatchNormalization提升收敛速度和泛化能力
多 GPU 训练使用 MirroredStrategy 加速训练
自定义损失函数如 Focal Loss、Label Smoothing
模型推理加载模型后进行预测或部署

U-Net

U-Net 是一种常用于医学图像分割的卷积神经网络模型。它特别适合于小数据集训练,通过使用数据增强、加权损失函数和高效的下采样与上采样路径来实现精确的像素级分类。

下面是一个基于 TensorFlow 通用训练模板实现 U-Net 模型,并以一个假设的数据集为例进行训练的完整代码示例。我们将不具体指定数据集,但会提供一个框架,您可以根据自己的数据集(如医学影像分割任务)调整输入部分。

使用 TensorFlow 通用模板实现 U-Net

import tensorflow as tf
from tensorflow.keras import layers, models, optimizers, callbacks, losses
import numpy as np
import os
import logging
import time


# ================== Step 1: 数据准备(这里以随机生成的数据为例)==================
def generate_random_dataset(num_samples=1000, img_shape=(128, 128, 1), mask_shape=(128, 128, 1)):
    """
    随机生成数据集供演示使用。
    """
    images = np.random.rand(num_samples, *img_shape)
    masks = np.random.randint(0, 2, size=(num_samples, *mask_shape))
    return images, masks

def get_dataset(batch_size=32):
    """
    加载并预处理数据集
    """
    images, masks = generate_random_dataset()

    dataset = tf.data.Dataset.from_tensor_slices((images, masks))
    dataset = dataset.shuffle(buffer_size=1024).batch(batch_size).prefetch(tf.data.AUTOTUNE)

    # 分割为训练集和验证集
    train_size = int(0.8 * len(images))
    train_dataset = dataset.take(train_size)
    val_dataset = dataset.skip(train_size)

    return train_dataset, val_dataset


# ================== Step 2: 构建 U-Net 模型 ==================
def conv_block(input_tensor, num_filters):
    x = layers.Conv2D(num_filters, (3, 3), padding='same')(input_tensor)
    x = layers.BatchNormalization()(x)
    x = layers.Activation('relu')(x)
    x = layers.Conv2D(num_filters, (3, 3), padding='same')(x)
    x = layers.BatchNormalization()(x)
    x = layers.Activation('relu')(x)
    return x

def encoder_block(input_tensor, num_filters):
    x = conv_block(input_tensor, num_filters)
    p = layers.MaxPooling2D((2, 2))(x)
    return x, p

def decoder_block(input_tensor, skip_features, num_filters):
    x = layers.Conv2DTranspose(num_filters, (2, 2), strides=2, padding='same')(input_tensor)
    x = layers.concatenate([x, skip_features])
    x = conv_block(x, num_filters)
    return x

def build_unet(img_shape):
    inputs = layers.Input(shape=img_shape)

    s1, p1 = encoder_block(inputs, 64)
    s2, p2 = encoder_block(p1, 128)
    s3, p3 = encoder_block(p2, 256)
    s4, p4 = encoder_block(p3, 512)

    b = conv_block(p4, 1024)

    d1 = decoder_block(b, s4, 512)
    d2 = decoder_block(d1, s3, 256)
    d3 = decoder_block(d2, s2, 128)
    d4 = decoder_block(d3, s1, 64)

    outputs = layers.Conv2D(1, (1, 1), activation='sigmoid')(d4)

    model = models.Model(inputs=[inputs], outputs=[outputs])
    return model


# ================== Step 3: 工具函数与配置 ==================
def setup_logger(log_dir="logs"):
    os.makedirs(log_dir, exist_ok=True)
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(levelname)s - %(message)s',
        handlers=[
            logging.FileHandler(f"{log_dir}/train_{time.strftime('%Y%m%d_%H%M%S')}.log"),
            logging.StreamHandler()
        ]
    )


# ================== Step 4: 主程序入口 ==================
def main():
    setup_logger()
    strategy = tf.distribute.MirroredStrategy()
    logging.info(f'Number of devices: {strategy.num_replicas_in_sync}')

    with strategy.scope():
        batch_size = 32
        epochs = 50
        img_shape = (128, 128, 1)

        # 获取数据集
        train_dataset, val_dataset = get_dataset(batch_size)

        # 构建模型
        model = build_unet(img_shape)

        # 编译模型
        optimizer = optimizers.Adam(learning_rate=1e-4)
        loss_fn = losses.BinaryCrossentropy(from_logits=False)
        metric = 'accuracy'

        model.compile(optimizer=optimizer,
                      loss=loss_fn,
                      metrics=[metric])

        # 回调函数
        checkpoint_cb = callbacks.ModelCheckpoint(
            "checkpoints/best_unet.keras",
            monitor='val_accuracy',
            save_best_only=True,
            mode='max'
        )
        early_stop_cb = callbacks.EarlyStopping(patience=10, restore_best_weights=True)
        reduce_lr = callbacks.ReduceLROnPlateau(monitor='val_loss', factor=0.2, patience=5)
        tensorboard_cb = callbacks.TensorBoard(log_dir="runs")

        callbacks_list = [checkpoint_cb, early_stop_cb, reduce_lr, tensorboard_cb]

        # 开始训练
        history = model.fit(
            train_dataset,
            validation_data=val_dataset,
            epochs=epochs,
            callbacks=callbacks_list
        )

        logging.info("Training completed.")
        model.save("checkpoints/final_unet.keras")


if __name__ == "__main__":
    main()

输出说明

文件夹内容
logs/存放训练日志文件
runs/TensorBoard 日志,用于可视化训练过程
checkpoints/保存最佳模型 best_unet.keras 和最终模型 final_unet.keras

关键说明

模型结构说明(U-Net)

  • 编码器(Encoder):通过一系列的卷积和最大池化操作提取特征图。每个块包括两个卷积层,接着是批量归一化和 ReLU 激活函数。
  • 解码器(Decoder):通过转置卷积操作逐步恢复空间分辨率。每个块接收来自编码器对应阶段的特征映射作为跳过连接,以帮助保留位置信息。
  • 输出层:应用一个卷积层将通道数减少到所需的类别数量(本例中为 1),并使用 Sigmoid 激活函数进行二分类。

数据预处理

  • 示例中使用了随机生成的数据集。在实际应用中,您需要替换这部分代码以加载和预处理您的特定数据集,例如医学影像数据。

模型编译

  • 使用 Adam 优化器,学习率默认为 1e-4
  • 损失函数为 BinaryCrossentropy(适用于二分类问题)
  • 指标为准确率

训练策略

  • 使用了 ModelCheckpoint 保存最佳模型
  • 使用了 ReduceLROnPlateau 动态调整学习率
  • 使用了 EarlyStopping 防止过拟合
  • 使用了 TensorBoard 可视化训练过程

扩展建议

功能描述
增加图像增强添加 RandomFlipRandomRotation 等数据增强操作
替换为真实数据集如医学影像数据集等
多 GPU 训练使用 MirroredStrategy 加速训练
自定义损失函数如 Dice Loss 或 Jaccard Index 等更适合分割任务的损失函数
模型推理加载模型后进行预测或部署

GANs (DCGAN)

下面是一个基于 TensorFlow 通用训练模板实现的经典 GAN(生成对抗网络) 的完整代码示例。我们将使用 DCGAN(深度卷积 GAN) 架构,在 MNIST 手写数字数据集 上进行训练,以演示如何从随机噪声生成逼真的图像。

使用 TensorFlow 通用模板实现 DCGAN

import tensorflow as tf
from tensorflow.keras import layers, models, optimizers, callbacks
from tensorflow.keras.datasets import mnist
import numpy as np
import os
import logging
import time
import matplotlib.pyplot as plt


# ================== Step 1: 数据准备 ==================
def get_dataset(batch_size=128):
    """
    加载并预处理 MNIST 数据集
    """
    (train_images, _), _ = mnist.load_data()
    train_images = train_images.reshape(-1, 28, 28, 1).astype("float32") / 127.5 - 1.0  # [-1, 1] 归一化

    dataset = tf.data.Dataset.from_tensor_slices(train_images)
    dataset = dataset.shuffle(buffer_size=60000).batch(batch_size).prefetch(tf.data.AUTOTUNE)

    return dataset


# ================== Step 2: 构建生成器和判别器 ==================

def build_generator(latent_dim):
    """
    生成器:将随机噪声向量映射为 28x28x1 图像
    """
    model = models.Sequential([
        layers.Dense(7 * 7 * 256, use_bias=False, input_shape=(latent_dim,)),
        layers.BatchNormalization(),
        layers.ReLU(),

        layers.Reshape((7, 7, 256)),

        layers.Conv2DTranspose(128, (5, 5), strides=(1, 1), padding='same', use_bias=False),
        layers.BatchNormalization(),
        layers.ReLU(),

        layers.Conv2DTranspose(64, (5, 5), strides=(2, 2), padding='same', use_bias=False),
        layers.BatchNormalization(),
        layers.ReLU(),

        layers.Conv2DTranspose(1, (5, 5), strides=(2, 2), padding='same', use_bias=False, activation='tanh')
    ])

    return model


def build_discriminator():
    """
    判别器:判断图像是真实的还是生成的
    """
    model = models.Sequential([
        layers.Input(shape=(28, 28, 1)),

        layers.Conv2D(64, (5, 5), strides=(2, 2), padding='same'),
        layers.LeakyReLU(alpha=0.2),

        layers.Conv2D(128, (5, 5), strides=(2, 2), padding='same'),
        layers.LeakyReLU(alpha=0.2),

        layers.Flatten(),
        layers.Dense(1)  # 输出一个实数分数(不加 sigmoid)
    ])

    return model


# ================== Step 3: 定义 GAN 模型 ==================
class DCGAN(models.Model):
    def __init__(self, generator, discriminator, latent_dim, **kwargs):
        super(DCGAN, self).__init__(**kwargs)
        self.generator = generator
        self.discriminator = discriminator
        self.latent_dim = latent_dim

    def compile(self, g_optimizer, d_optimizer, loss_fn):
        super(DCGAN, self).compile()
        self.g_optimizer = g_optimizer
        self.d_optimizer = d_optimizer
        self.loss_fn = loss_fn

    def train_step(self, real_images):

        batch_size = tf.shape(real_images)[0]

        # ----------------------------
        # 训练判别器
        # ----------------------------
        noise = tf.random.normal(shape=(batch_size, self.latent_dim))
        generated_images = self.generator(noise)

        with tf.GradientTape() as tape:
            real_output = self.discriminator(real_images, training=True)
            fake_output = self.discriminator(generated_images, training=True)

            real_loss = self.loss_fn(tf.ones_like(real_output), real_output)
            fake_loss = self.loss_fn(tf.zeros_like(fake_output), fake_output)
            d_loss = real_loss + fake_loss

        grads = tape.gradient(d_loss, self.discriminator.trainable_weights)
        self.d_optimizer.apply_gradients(zip(grads, self.discriminator.trainable_weights))

        # ----------------------------
        # 训练生成器
        # ----------------------------
        noise = tf.random.normal(shape=(batch_size, self.latent_dim))

        with tf.GradientTape() as tape:
            generated_images = self.generator(noise, training=True)
            fake_output = self.discriminator(generated_images, training=False)

            g_loss = self.loss_fn(tf.ones_like(fake_output), fake_output)

        grads = tape.gradient(g_loss, self.generator.trainable_weights)
        self.g_optimizer.apply_gradients(zip(grads, self.generator.trainable_weights))

        return {"d_loss": d_loss, "g_loss": g_loss}


# ================== Step 4: 工具函数与配置 ==================
def setup_logger(log_dir="logs"):
    os.makedirs(log_dir, exist_ok=True)
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(levelname)s - %(message)s',
        handlers=[
            logging.FileHandler(f"{log_dir}/train_{time.strftime('%Y%m%d_%H%M%S')}.log"),
            logging.StreamHandler()
        ]
    )

def plot_generated_images(epoch, generator, test_input, save_path="images"):
    predictions = generator(test_input, training=False)
    fig, axes = plt.subplots(4, 4, figsize=(8, 8))
    for i, ax in enumerate(axes.flat):
        ax.imshow(predictions[i, :, :, 0] * 0.5 + 0.5, cmap='gray')
        ax.axis("off")
    plt.tight_layout()
    os.makedirs(save_path, exist_ok=True)
    plt.savefig(f"{save_path}/image_at_epoch_{epoch:04d}.png")
    plt.close()


# ================== Step 5: 主程序入口 ==================
def main():
    setup_logger()
    strategy = tf.distribute.MirroredStrategy()
    logging.info(f'Number of devices: {strategy.num_replicas_in_sync}')

    with strategy.scope():
        batch_size = 128
        epochs = 100
        latent_dim = 100

        # 获取数据集
        dataset = get_dataset(batch_size)

        # 构建模型
        generator = build_generator(latent_dim)
        discriminator = build_discriminator()

        # 编译模型
        dcgan = DCGAN(generator, discriminator, latent_dim)
        dcgan.compile(
            g_optimizer=optimizers.Adam(learning_rate=2e-4, beta_1=0.5),
            d_optimizer=optimizers.Adam(learning_rate=2e-4, beta_1=0.5),
            loss_fn=losses.Hinge()  # 可选 Hinge Loss 或 BinaryCrossentropy
        )

        # 回调函数
        checkpoint_cb = callbacks.ModelCheckpoint(
            "checkpoints/generator.keras",
            monitor='g_loss',
            save_best_only=True,
            mode='min'
        )
        tensorboard_cb = callbacks.TensorBoard(log_dir="runs")

        fixed_noise = tf.random.normal([16, latent_dim])  # 固定噪声用于可视化生成效果

        # 自定义回调:定期保存生成图像
        class ImageSaveCallback(callbacks.Callback):
            def on_epoch_end(self, epoch, logs=None):
                if (epoch + 1) % 5 == 0:
                    plot_generated_images(epoch + 1, self.model.generator, fixed_noise)

        image_save_cb = ImageSaveCallback()

        callbacks_list = [checkpoint_cb, tensorboard_cb, image_save_cb]

        # 开始训练
        history = dcgan.fit(
            dataset,
            epochs=epochs,
            callbacks=callbacks_list
        )

        logging.info("Training completed.")
        dcgan.generator.save("checkpoints/final_generator.keras")
        dcgan.discriminator.save("checkpoints/final_discriminator.keras")


if __name__ == "__main__":
    main()

输出说明

文件夹内容
logs/存放训练日志文件
runs/TensorBoard 日志,用于可视化训练过程
images/每隔一定轮次保存的生成图像
checkpoints/保存最佳生成器、最终生成器和判别器

关键说明

模型结构说明(DCGAN)

部分类型参数说明
GeneratorConv2DTranspose将噪声向量映射为真实图像
DiscriminatorConv2D判断图像是真实还是生成的
Loss FunctionHinge Loss更适合稳定训练 GAN
OptimizerAdam(2e-4, beta=0.5)常用于 GAN 的优化器参数

数据预处理

  • 使用了标准的 MNIST 数据集
  • 图像大小为 (28, 28, 1),归一化到 [-1, 1]
  • 使用 tf.data.Dataset 进行高效加载和批处理

模型编译与训练

  • 使用自定义 train_step() 来实现交替训练策略
  • 判别器使用真实标签 [1] 和生成标签 [0] 来计算损失
  • 生成器的目标是让判别器输出接近 [1]

扩展建议

功能描述
替换为其他架构如 CycleGAN、StyleGAN、Wasserstein-GAN 等
多 GPU 训练使用 MirroredStrategy 加速训练
自定义损失函数如 Wasserstein loss、LS-GAN loss
图像转译任务如 pix2pix、CycleGAN 等条件 GAN
模型推理用固定噪声生成图像或图像修复等

Transformer

实现一个基于 TensorFlow 的 Transformer 模型涉及多个步骤,包括定义模型架构、准备数据集、设置训练循环以及评估模型。下面是一个简化的示例,展示了如何使用 TensorFlow 实现一个基本的 Transformer 模型进行序列到序列(seq2seq)任务,比如机器翻译。

代码示例

import tensorflow as tf
from tensorflow.keras import layers

# Transformer 参数
num_layers = 4
d_model = 128
dff = 512
num_heads = 8
dropout_rate = 0.1

def get_angles(pos, i, d_model):
    angle_rates = 1 / np.power(10000, (2 * (i // 2)) / np.float32(d_model))
    return pos * angle_rates

def positional_encoding(position, d_model):
    angle_rads = get_angles(np.arange(position)[:, np.newaxis],
                            np.arange(d_model)[np.newaxis, :],
                            d_model)
    
    # 将 sin 应用于数组中的偶数索引:2i
    angle_rads[:, 0::2] = np.sin(angle_rads[:, 0::2])
    
    # 将 cos 应用于数组中的奇数索引:2i+1
    angle_rads[:, 1::2] = np.cos(angle_rads[:, 1::2])
    
    pos_encoding = angle_rads[np.newaxis, ...]
    
    return tf.cast(pos_encoding, dtype=tf.float32)

class MultiHeadAttention(layers.Layer):
    def __init__(self, d_model, num_heads):
        super(MultiHeadAttention, self).__init__()
        self.num_heads = num_heads
        self.d_model = d_model
        
        assert d_model % self.num_heads == 0
        
        self.depth = d_model // self.num_heads
        
        self.wq = layers.Dense(d_model)
        self.wk = layers.Dense(d_model)
        self.wv = layers.Dense(d_model)
        
        self.dense = layers.Dense(d_model)
    
    # 省略了 call 方法的具体实现,它应包括 Q, K, V 计算和多头机制

class EncoderLayer(layers.Layer):
    def __init__(self, d_model, num_heads, dff, rate=0.1):
        super(EncoderLayer, self).__init__()

        self.mha = MultiHeadAttention(d_model, num_heads)
        self.ffn = tf.keras.Sequential([
            layers.Dense(dff, activation='relu'),  # (batch_size, seq_len, dff)
            layers.Dense(d_model)  # (batch_size, seq_len, d_model)
        ])

        self.layernorm1 = layers.LayerNormalization(epsilon=1e-6)
        self.layernorm2 = layers.LayerNormalization(epsilon=1e-6)

        self.dropout1 = layers.Dropout(rate)
        self.dropout2 = layers.Dropout(rate)
    
    # 省略了 call 方法的具体实现

class Encoder(layers.Layer):
    def __init__(self, num_layers, d_model, num_heads, dff,
                 maximum_position_encoding, rate=0.1):
        super(Encoder, self).__init__()

        self.d_model = d_model
        self.num_layers = num_layers
        
        self.embedding = layers.Embedding(input_dim=8500, output_dim=d_model)
        self.pos_encoding = positional_encoding(maximum_position_encoding, self.d_model)
        
        self.enc_layers = [EncoderLayer(d_model, num_heads, dff, rate) 
                           for _ in range(num_layers)]

        self.dropout = layers.Dropout(rate)
    
    # 省略了 call 方法的具体实现

# 构建 Transformer 模型
def build_transformer(vocab_size, num_layers, d_model, num_heads, dff, input_vocab_size,
                      target_vocab_size, pe_input, pe_target, rate=0.1):
    inputs = layers.Input(shape=(None,))
    targets = layers.Input(shape=(None,))

    enc_padding_mask = layers.Lambda(
        create_padding_mask, output_shape=(1, 1, None),
        name='enc_padding_mask')(inputs)
    
    encoder = Encoder(num_layers=num_layers, d_model=d_model,
                      num_heads=num_heads, dff=dff,
                      maximum_position_encoding=pe_input)

    enc_output = encoder(inputs, training=False, mask=enc_padding_mask)
    
    # 这里需要添加解码器和其他组件以完成Transformer
    
    transformer = tf.keras.Model(inputs=[inputs, targets], outputs=final_output)
    return transformer

# 示例:构建并编译模型
transformer = build_transformer(
    vocab_size=8500,
    num_layers=4,
    d_model=128,
    num_heads=8,
    dff=512,
    input_vocab_size=8500,
    target_vocab_size=8000,
    pe_input=10000,
    pe_target=6000)

transformer.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])

print(transformer.summary())

BERT

下面是一个基于 TensorFlow 通用训练模板实现的简化版 BERT(Bidirectional Encoder Representations from Transformers) 模型,用于文本分类任务。我们将使用 IMDB 影评数据集 进行二分类情感分析任务,并展示如何:

  • 构建 BERT 基础编码器结构
  • 使用 Tokenizer 和 Positional Encoding
  • 加载和预处理文本数据
  • 编写完整训练流程(包括模型编译、回调函数、日志记录等)

使用 TensorFlow 实现简化版 BERT(适用于文本分类)

注意:本示例为简化版 BERT(Mini-BERT),适合教学与快速实验,非 HuggingFace 提供的完整预训练 BERT 模型。

环境依赖

pip install tensorflow tensorflow-text

完整代码

import tensorflow as tf
from tensorflow.keras import layers, models, optimizers, callbacks, losses, metrics
import tensorflow_text as text
import numpy as np
import os
import logging
import time


# ================== Step 1: 数据准备 ==================
def get_dataset(batch_size=32):
    """
    加载 IMDB 数据集并进行简单预处理
    """
    (x_train, y_train), (x_test, y_test) = tf.keras.datasets.imdb.load_data(num_words=None)
    x_train = tf.keras.preprocessing.sequence.pad_sequences(x_train, maxlen=512, padding='post', truncating='post')
    x_test = tf.keras.preprocessing.sequence.pad_sequences(x_test, maxlen=512, padding='post', truncating='post')

    train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train)).shuffle(10000).batch(batch_size).prefetch(tf.data.AUTOTUNE)
    val_dataset = tf.data.Dataset.from_tensor_slices((x_test, y_test)).batch(batch_size).prefetch(tf.data.AUTOTUNE)

    return train_dataset, val_dataset


# ================== Step 2: 构建 Mini-BERT 模型 ==================
def gelu(x):
    """Gaussian Error Linear Unit activation function"""
    return 0.5 * x * (1 + tf.tanh(tf.sqrt(2 / np.pi) * (x + 0.044715 * tf.pow(x, 3))))


class TokenAndPositionEmbedding(layers.Layer):
    def __init__(self, vocab_size, max_len, embed_dim, **kwargs):
        super(TokenAndPositionEmbedding, self).__init__(**kwargs)
        self.token_emb = layers.Embedding(input_dim=vocab_size, output_dim=embed_dim, mask_zero=True)
        self.pos_emb = layers.Embedding(input_dim=max_len, output_dim=embed_dim)

    def call(self, x):
        max_len = tf.shape(x)[-1]
        positions = tf.range(start=0, limit=max_len, delta=1)
        position_embeddings = self.pos_emb(positions)
        token_embeddings = self.token_emb(x)
        return token_embeddings + position_embeddings


class TransformerBlock(layers.Layer):
    def __init__(self, embed_dim, num_heads, ff_dim, rate=0.1, **kwargs):
        super(TransformerBlock, self).__init__(**kwargs)
        self.att = layers.MultiHeadAttention(num_heads=num_heads, key_dim=embed_dim)
        self.ffn = tf.keras.Sequential([
            layers.Dense(ff_dim, activation=gelu),
            layers.Dense(embed_dim)
        ])
        self.layernorm1 = layers.LayerNormalization(epsilon=1e-6)
        self.layernorm2 = layers.LayerNormalization(epsilon=1e-6)
        self.dropout1 = layers.Dropout(rate)
        self.dropout2 = layers.Dropout(rate)

    def call(self, inputs, training=False):
        attn_output = self.att(inputs, inputs)
        attn_output = self.dropout1(attn_output, training=training)
        out1 = self.layernorm1(inputs + attn_output)

        ffn_output = self.ffn(out1)
        ffn_output = self.dropout2(ffn_output, training=training)
        return self.layernorm2(out1 + ffn_output)


def build_bert_model(vocab_size=10000, max_len=512, embed_dim=128, num_heads=8, ff_dim=512, num_transformer_blocks=4):
    inputs = layers.Input(shape=(max_len,), dtype=tf.int32)
    embedding_layer = TokenAndPositionEmbedding(vocab_size, max_len, embed_dim)
    x = embedding_layer(inputs)

    for _ in range(num_transformer_blocks):
        x = TransformerBlock(embed_dim, num_heads, ff_dim)(x)

    cls_token = layers.Lambda(lambda x: x[:, 0])(x)  # 取第一个 token [CLS] 作为句子表示
    outputs = layers.Dense(1, activation='sigmoid')(cls_token)

    model = models.Model(inputs=inputs, outputs=outputs)
    return model


# ================== Step 3: 工具函数与配置 ==================
def setup_logger(log_dir="logs"):
    os.makedirs(log_dir, exist_ok=True)
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(levelname)s - %(message)s',
        handlers=[
            logging.FileHandler(f"{log_dir}/train_{time.strftime('%Y%m%d_%H%M%S')}.log"),
            logging.StreamHandler()
        ]
    )


# ================== Step 4: 主程序入口 ==================
def main():
    setup_logger()
    strategy = tf.distribute.MirroredStrategy()
    logging.info(f'Number of devices: {strategy.num_replicas_in_sync}')

    with strategy.scope():
        batch_size = 32
        epochs = 10
        vocab_size = 10000
        max_len = 512

        # 获取数据集
        train_dataset, val_dataset = get_dataset(batch_size)

        # 构建模型
        model = build_bert_model(vocab_size=vocab_size, max_len=max_len)

        # 编译模型
        optimizer = optimizers.AdamW(learning_rate=3e-5, weight_decay=0.001)
        loss_fn = losses.BinaryCrossentropy()
        metric = metrics.BinaryAccuracy(name='acc')

        model.compile(optimizer=optimizer,
                      loss=loss_fn,
                      metrics=[metric])

        # 回调函数
        checkpoint_cb = callbacks.ModelCheckpoint(
            "checkpoints/best_bert.keras",
            monitor='val_acc',
            save_best_only=True,
            mode='max'
        )
        early_stop_cb = callbacks.EarlyStopping(patience=3, restore_best_weights=True)
        reduce_lr = callbacks.ReduceLROnPlateau(monitor='val_loss', factor=0.2, patience=2)
        tensorboard_cb = callbacks.TensorBoard(log_dir="runs")

        callbacks_list = [checkpoint_cb, early_stop_cb, reduce_lr, tensorboard_cb]

        # 开始训练
        history = model.fit(
            train_dataset,
            validation_data=val_dataset,
            epochs=epochs,
            callbacks=callbacks_list
        )

        logging.info("Training completed.")
        model.save("checkpoints/final_bert.keras")


if __name__ == "__main__":
    main()

输出说明

文件夹内容
logs/存放训练日志文件
runs/TensorBoard 日志,用于可视化训练过程
checkpoints/保存最佳模型 best_bert.keras 和最终模型 final_bert.keras

关键说明

模型结构说明(Mini-BERT)

组件描述
TokenAndPositionEmbedding同时学习词嵌入和位置编码
TransformerBlock包含多头注意力和前馈网络,采用 LayerNorm 和 GELU 激活
[CLS] token取第一个 token 作为句子表示用于分类
AdamW带权重衰减的优化器,适合 BERT 类模型

数据预处理

  • 使用了标准的 IMDB 数据集(影评情感分类)
  • 文本被转换为整数序列并填充到统一长度 max_len=512
  • 标签是二值标签(正面/负面)

模型编译

  • 使用 AdamW 优化器(推荐用于 BERT 类模型)
  • 损失函数为 BinaryCrossentropy(适用于二分类)
  • 指标为准确率

训练策略

  • 使用了 ModelCheckpoint 保存最佳模型
  • 使用了 ReduceLROnPlateau 动态调整学习率
  • 使用了 EarlyStopping 防止过拟合
  • 使用了 TensorBoard 可视化训练过程

扩展建议

功能描述
使用 HuggingFace Transformers直接加载 bert-base-uncased 等官方模型
多分类任务修改输出层为 Dense(num_classes) 并使用 SparseCategoricalCrossentropy
NER、问答等任务添加对应的任务头(Task Head)
自定义 Tokenizer使用 BertTokenizer 或 SentencePiece 分词器
多 GPU 训练使用 MirroredStrategy 加速训练
模型推理加载模型后进行预测或部署

AutoRec 推荐

简介:AutoRec 推荐模型

  • 提出者:Quadratic Probing Lab, NUS
  • 年份:2015
  • 核心思想
    • 将协同过滤问题转化为矩阵补全任务
    • 使用自编码器(Autoencoder)对用户/物品评分向量进行编码与重构
    • 用户隐向量或物品隐向量作为推荐依据
  • 两种变体
    • User-based AutoRec:输入是用户的评分向量(所有物品),输出是该用户的预测评分
    • Item-based AutoRec:输入是物品的评分向量(所有用户),输出是该物品的预测评分
  • 优点
    • 结构简单,适合冷启动和稀疏评分数据
    • 可以通过正则化防止过拟合
    • 易于扩展为深度结构

模型结构简述(以 Item-based AutoRec 为例)

组件描述
输入层物品的评分向量(长度 = 用户数)
编码层Dense 层压缩为低维表示
解码层Dense 层还原为原始维度
输出层Reconstructed 评分向量
损失函数MSE Loss(仅计算已评分项)

使用 TensorFlow/Keras 模板完成 AutoRec 实现

以下是基于你给的模板构建的 Item-based AutoRec 推荐模型:

import tensorflow as tf
from tensorflow.keras import layers, models, optimizers, callbacks, losses, metrics
import os
import logging
import time
import numpy as np


# ================== Step 1: 数据准备 ==================
def get_dataset(num_users=1000, num_items=500, sparsity_ratio=0.95):
    """
    随机生成一个稀疏的评分矩阵(可替换为真实数据)
    :param num_users: 用户数量
    :param num_items: 物品数量
    :param sparsity_ratio: 稀疏率(0~1)
    :return: 训练集、验证集(物品视角)
    """
    ratings = np.random.randint(1, 6, size=(num_users, num_items)).astype(np.float32)
    mask = np.random.rand(*ratings.shape) < sparsity_ratio
    ratings[mask] = 0  # 用 0 表示未评分

    # 转换为物品视角(每个样本是一个物品的评分向量)
    item_ratings = ratings.T  # shape: (num_items, num_users)

    # 构建 Dataset(仅训练非零评分)
    train_data = []
    val_mask = []

    for i in range(num_items):
        vec = item_ratings[i]
        zero_mask = (vec == 0)
        if np.sum(~zero_mask) == 0:
            continue  # 跳过无评分的物品
        train_data.append(vec)
        val_mask.append(zero_mask)

    train_data = np.array(train_data)
    val_mask = np.array(val_mask)

    return train_data, val_mask


# ================== Step 2: 构建模型 ==================
def build_model(input_dim, hidden_units=[256, 128]):
    """
    构建 Item-based AutoRec 模型
    :param input_dim: 输入维度(用户数量)
    :param hidden_units: 编码器和解码器的隐藏层大小
    """
    inputs = tf.keras.Input(shape=(input_dim,))
    x = layers.Dense(hidden_units[0], activation='relu')(inputs)
    encoded = layers.Dense(hidden_units[1], activation='relu')(x)
    x = layers.Dense(hidden_units[0], activation='relu')(encoded)
    outputs = layers.Dense(input_dim)(x)

    model = models.Model(inputs=inputs, outputs=outputs)

    return model


# ================== Step 3: 工具函数与配置 ==================
def setup_logger(log_dir="logs"):
    os.makedirs(log_dir, exist_ok=True)
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(levelname)s - %(message)s',
        handlers=[
            logging.FileHandler(f"{log_dir}/train_{time.strftime('%Y%m%d_%H%M%S')}.log"),
            logging.StreamHandler()
        ]
    )


# ================== Step 4: 主程序入口 ==================
def main():
    setup_logger()
    strategy = tf.distribute.MirroredStrategy()
    logging.info(f'Number of devices: {strategy.num_replicas_in_sync}')

    with strategy.scope():
        num_users = 1000
        num_items = 500
        batch_size = 128
        epochs = 50
        input_dim = num_users  # 每个物品的输入是用户对该物品的评分向量
        hidden_units = [256, 128]

        # 获取数据集(替换为你自己的 Dataset)
        train_data, val_mask = get_dataset(num_users=num_users, num_items=num_items)
        dataset = tf.data.Dataset.from_tensor_slices((train_data))
        dataset = dataset.shuffle(buffer_size=1000).batch(batch_size).prefetch(tf.data.AUTOTUNE)

        # 构建模型
        model = build_model(input_dim, hidden_units)

        # 自定义损失函数:只计算非零评分部分的 MSE
        def masked_mse(y_true, y_pred):
            mask = tf.not_equal(y_true, 0)
            masked_true = tf.boolean_mask(y_true, mask)
            masked_pred = tf.boolean_mask(y_pred, mask)
            return tf.reduce_mean(tf.square(masked_true - masked_pred))

        # 编译模型
        optimizer = optimizers.Adam(learning_rate=1e-3)
        model.compile(optimizer=optimizer, loss=masked_mse)

        # 回调函数
        checkpoint_cb = callbacks.ModelCheckpoint(
            "checkpoints/best_model.keras",
            monitor='loss',
            save_best_only=True,
            mode='min'
        )
        early_stop_cb = callbacks.EarlyStopping(patience=10, restore_best_weights=True)
        tensorboard_cb = callbacks.TensorBoard(log_dir="runs")

        callbacks_list = [checkpoint_cb, early_stop_cb, tensorboard_cb]

        # 开始训练
        history = model.fit(
            dataset,
            epochs=epochs,
            callbacks=callbacks_list
        )

        logging.info("Training completed.")
        model.save("checkpoints/final_model.keras")

        # 示例:预测并评估
        pred_ratings = model.predict(train_data)
        for idx in range(5):  # 打印前5个物品的预测结果
            logging.info(f"Item {idx}: Predictions: {pred_ratings[idx][val_mask[idx][:]]}")


if __name__ == "__main__":
    main()

总结

本代码实现了 Item-based AutoRec 推荐模型,适用于评分矩阵补全任务。你可以根据实际业务需求进一步优化模型结构、引入负采样、添加时间因素、迁移到真实数据集等。


Neural Collaborative Filtering (NCF)

简介:Neural Collaborative Filtering(NCF)

  • 提出者:新加坡国立大学 & 谷歌
  • 年份:2017
  • 核心思想
    • 将传统的矩阵分解推广到深度神经网络中
    • 使用 广义矩阵分解(GMF) + 多层感知机(MLP) 的混合结构
    • 同时建模线性和非线性用户-物品交互关系
  • 两种主要组件
    • GMF(Generalized Matrix Factorization):模拟用户和物品的隐向量内积
    • MLP(Multi-Layer Perceptron):建模复杂的用户-物品特征组合
  • 输出层:将 GMF 和 MLP 的输出拼接后通过全连接层进行预测

模型结构简述

组件描述
输入用户 ID、物品 ID
Embedding 层分别将用户和物品映射为低维隐向量
GMF 部分用户和物品嵌入的 element-wise product(或 inner product)
MLP 部分用户和物品嵌入的拼接 + 全连接层
输出层GMF + MLP 输出拼接并通过 Dense 层得到评分预测
损失函数Binary Crossentropy(用于隐式反馈)或 MSE(用于显式评分)

使用 TensorFlow/Keras 模板完成 NCF 实现

以下是基于你给的模板构建的 NeuMF(NCF 中的混合版本) 推荐模型:

import tensorflow as tf
from tensorflow.keras import layers, models, optimizers, callbacks, losses, metrics
import os
import logging
import time
import numpy as np


# ================== Step 1: 数据准备 ==================
def get_dataset(num_users=1000, num_items=500, num_interactions=10000, test_size=0.2):
    """
    构造一个隐式反馈数据集(用户-物品点击/浏览)
    :param num_users: 用户数量
    :param num_items: 物品数量
    :param num_interactions: 总交互次数
    :param test_size: 测试集比例
    """
    interactions = []
    for _ in range(num_interactions):
        user_id = np.random.randint(0, num_users)
        item_id = np.random.randint(0, num_items)
        interactions.append((user_id, item_id))

    # 划分训练集和验证集
    np.random.shuffle(interactions)
    split_idx = int(len(interactions) * (1 - test_size))
    train_pairs = interactions[:split_idx]
    val_pairs = interactions[split_idx:]

    return train_pairs, val_pairs, num_users, num_items


# ================== Step 2: 构建模型 ==================
class NCFModel(models.Model):
    def __init__(self, num_users, num_items, embedding_dim=64, mlp_layers=[64, 32, 16], **kwargs):
        super(NCFModel, self).__init__(**kwargs)
        self.num_users = num_users
        self.num_items = num_items
        self.embedding_dim = embedding_dim

        # 用户和物品的嵌入层
        self.user_embedding_gmf = layers.Embedding(
            input_dim=num_users, output_dim=embedding_dim, name='user_embedding_gmf')
        self.item_embedding_gmf = layers.Embedding(
            input_dim=num_items, output_dim=embedding_dim, name='item_embedding_gmf')

        self.user_embedding_mlp = layers.Embedding(
            input_dim=num_users, output_dim=embedding_dim, name='user_embedding_mlp')
        self.item_embedding_mlp = layers.Embedding(
            input_dim=num_items, output_dim=embedding_dim, name='item_embedding_mlp')

        # MLP 部分
        mlp_layers_list = []
        for units in mlp_layers:
            mlp_layers_list.append(layers.Dense(units, activation='relu'))
        self.mlp = models.Sequential(mlp_layers_list, name='mlp')

        # 输出层
        self.final_dense = layers.Dense(1, activation='sigmoid', name='output')

    def call(self, inputs):
        user_ids, item_ids = inputs

        # GMF 部分
        user_gmf = self.user_embedding_gmf(user_ids)
        item_gmf = self.item_embedding_gmf(item_ids)
        gmf_vector = tf.multiply(user_gmf, item_gmf)  # Element-wise 乘积

        # MLP 部分
        user_mlp = self.user_embedding_mlp(user_ids)
        item_mlp = self.item_embedding_mlp(item_ids)
        mlp_vector = tf.concat([user_mlp, item_mlp], axis=-1)
        mlp_vector = self.mlp(mlp_vector)

        # 拼接 GMF 和 MLP 输出
        concat_vector = tf.concat([gmf_vector, mlp_vector], axis=-1)
        logits = self.final_dense(concat_vector)

        return logits


# ================== Step 3: 工具函数与配置 ==================
def setup_logger(log_dir="logs"):
    os.makedirs(log_dir, exist_ok=True)
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(levelname)s - %(message)s',
        handlers=[
            logging.FileHandler(f"{log_dir}/train_{time.strftime('%Y%m%d_%H%M%S')}.log"),
            logging.StreamHandler()
        ]
    )


# ================== Step 4: 主程序入口 ==================
def main():
    setup_logger()
    strategy = tf.distribute.MirroredStrategy()
    logging.info(f'Number of devices: {strategy.num_replicas_in_sync}')

    with strategy.scope():
        num_users = 1000
        num_items = 500
        batch_size = 256
        epochs = 50
        embedding_dim = 64
        mlp_layers = [64, 32, 16]

        # 获取数据集(替换为你自己的 Dataset)
        train_pairs, val_pairs, num_users, num_items = get_dataset(num_users, num_items)

        # 构造 tf.data.Dataset
        def make_dataset(pairs):
            users, items = zip(*pairs)
            x = {"user_input": np.array(users), "item_input": np.array(items)}
            y = np.ones(len(pairs))  # 假设所有样本都是正样本(隐式反馈)
            return tf.data.Dataset.from_tensor_slices((x, y)).shuffle(1000).batch(batch_size).prefetch(tf.data.AUTOTUNE)

        train_dataset = make_dataset(train_pairs)
        val_dataset = make_dataset(val_pairs)

        # 构建模型
        model = NCFModel(num_users, num_items, embedding_dim=embedding_dim, mlp_layers=mlp_layers)

        # 编译模型
        optimizer = optimizers.Adam(learning_rate=1e-3)
        loss_fn = losses.BinaryCrossentropy()
        metric = metrics.BinaryAccuracy(name='acc')

        model.compile(optimizer=optimizer, loss=loss_fn, metrics=[metric])

        # 回调函数
        checkpoint_cb = callbacks.ModelCheckpoint(
            "checkpoints/best_model.keras",
            monitor='val_loss',
            save_best_only=True,
            mode='min'
        )
        early_stop_cb = callbacks.EarlyStopping(patience=10, restore_best_weights=True)
        tensorboard_cb = callbacks.TensorBoard(log_dir="runs")

        callbacks_list = [checkpoint_cb, early_stop_cb, tensorboard_cb]

        # 开始训练
        history = model.fit(
            train_dataset,
            validation_data=val_dataset,
            epochs=epochs,
            callbacks=callbacks_list
        )

        logging.info("Training completed.")
        model.save("checkpoints/final_model.keras")


if __name__ == "__main__":
    main()

总结

本代码实现了 NeuMF(NCF) 推荐模型,适用于隐式反馈下的推荐任务。你可以根据实际业务需求进一步优化模型结构、加入负采样、使用更复杂网络、迁移到真实数据集等。


Wide & Deep

简介:Wide & Deep 推荐模型

  • 提出者:Google Research
  • 年份:2016
  • 核心思想
    • 同时建模 宽线性模型(Wide)深度神经网络(Deep)
    • Wide 部分用于记忆用户和物品的强特征组合(显式规则)
    • Deep 部分用于泛化,学习稀疏特征之间的潜在关系
  • 典型应用场景
    • 推荐系统(如 Google Play App 推荐)
    • 广告点击率预估(CTR)
    • 搜索排序优化
  • 输入结构
    • wide_input:稠密或 one-hot 编码的交叉特征(适合线性模型)
    • deep_input:嵌入后的稀疏特征(适合 DNN)

模型结构简述

组件描述
Wide 部分线性回归层,输入为交叉特征(如用户×物品)
Deep 部分多层感知机(MLP),输入为稀疏特征嵌入拼接
输出层将 Wide 和 Deep 的输出相加并通过 sigmoid 得到预测概率
损失函数Binary Crossentropy(隐式反馈)

使用 TensorFlow/Keras 模板完成 Wide & Deep 实现

以下是基于你给的模板构建的 Wide & Deep 推荐模型:

import tensorflow as tf
from tensorflow.keras import layers, models, optimizers, callbacks, losses, metrics
import os
import logging
import time
import numpy as np


# ================== Step 1: 数据准备 ==================
def get_dataset(num_samples=10000, num_users=1000, num_items=500, test_size=0.2):
    """
    构造一个模拟的推荐数据集(用户 ID + 物品 ID + 标签)
    :param num_samples: 总样本数
    :param num_users: 用户数量
    :param num_items: 物品数量
    :param test_size: 测试集比例
    """
    users = np.random.randint(0, num_users, size=num_samples)
    items = np.random.randint(0, num_items, size=num_samples)
    labels = np.random.randint(0, 2, size=num_samples)  # 二分类任务(点击/未点击)

    indices = np.random.permutation(num_samples)
    split_idx = int((1 - test_size) * num_samples)
    train_idx, val_idx = indices[:split_idx], indices[split_idx:]

    train_data = (users[train_idx], items[train_idx])
    train_labels = labels[train_idx]
    val_data = (users[val_idx], items[val_idx])
    val_labels = labels[val_idx]

    return (train_data, train_labels), (val_data, val_labels), num_users, num_items


# ================== Step 2: 构建模型 ==================
class WideAndDeepModel(models.Model):
    def __init__(self, num_users, num_items, embedding_dim=32, hidden_units=[64, 32], **kwargs):
        super(WideAndDeepModel, self).__init__(**kwargs)

        # Wide部分:用户 × 物品 交叉特征(one-hot)
        self.wide_dense = layers.Dense(1)

        # Deep部分:用户和物品嵌入 + MLP
        self.user_embedding = layers.Embedding(input_dim=num_users, output_dim=embedding_dim)
        self.item_embedding = layers.Embedding(input_dim=num_items, output_dim=embedding_dim)

        self.deep_mlp = models.Sequential([
            layers.Dense(units, activation='relu') for units in hidden_units
        ], name="deep_mlp")

        self.final_dense = layers.Dense(1, activation='sigmoid')

    def call(self, inputs):
        user_ids, item_ids = inputs

        # Wide 部分:用户 ID 和物品 ID 的 one-hot 交叉特征
        wide_input = tf.stack([user_ids, item_ids], axis=1)
        wide_one_hot = tf.one_hot(wide_input, depth=[num_users, num_items])  # shape: (batch_size, 2)
        wide_logits = self.wide_dense(tf.concat(wide_one_hot, axis=1))

        # Deep 部分:嵌入拼接 + MLP
        user_emb = self.user_embedding(user_ids)  # shape: (batch_size, embedding_dim)
        item_emb = self.item_embedding(item_ids)  # shape: (batch_size, embedding_dim)
        deep_input = tf.concat([user_emb, item_emb], axis=1)
        deep_logits = self.deep_mlp(deep_input)

        # 合并 Wide 和 Deep 的输出
        logits = self.final_dense(tf.concat([wide_logits, deep_logits], axis=1))
        return logits


# ================== Step 3: 工具函数与配置 ==================
def setup_logger(log_dir="logs"):
    os.makedirs(log_dir, exist_ok=True)
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(levelname)s - %(message)s',
        handlers=[
            logging.FileHandler(f"{log_dir}/train_{time.strftime('%Y%m%d_%H%M%S')}.log"),
            logging.StreamHandler()
        ]
    )


# ================== Step 4: 主程序入口 ==================
def main():
    setup_logger()
    strategy = tf.distribute.MirroredStrategy()
    logging.info(f'Number of devices: {strategy.num_replicas_in_sync}')

    with strategy.scope():
        num_users = 1000
        num_items = 500
        batch_size = 512
        epochs = 50
        embedding_dim = 32
        hidden_units = [64, 32]

        # 获取数据集(替换为你自己的 Dataset)
        (train_data, train_labels), (val_data, val_labels), _, _ = get_dataset(num_users=num_users, num_items=num_items)
        train_dataset = tf.data.Dataset.from_tensor_slices((train_data, train_labels)).shuffle(1000).batch(batch_size).prefetch(tf.data.AUTOTUNE)
        val_dataset = tf.data.Dataset.from_tensor_slices((val_data, val_labels)).batch(batch_size).prefetch(tf.data.AUTOTUNE)

        # 构建模型
        model = WideAndDeepModel(num_users, num_items, embedding_dim=embedding_dim, hidden_units=hidden_units)

        # 编译模型
        optimizer = optimizers.Adam(learning_rate=1e-3)
        loss_fn = losses.BinaryCrossentropy()
        metric = metrics.BinaryAccuracy(name='acc')

        model.compile(optimizer=optimizer, loss=loss_fn, metrics=[metric])

        # 回调函数
        checkpoint_cb = callbacks.ModelCheckpoint(
            "checkpoints/best_model.keras",
            monitor='val_loss',
            save_best_only=True,
            mode='min'
        )
        early_stop_cb = callbacks.EarlyStopping(patience=10, restore_best_weights=True)
        tensorboard_cb = callbacks.TensorBoard(log_dir="runs")

        callbacks_list = [checkpoint_cb, early_stop_cb, tensorboard_cb]

        # 开始训练
        history = model.fit(
            train_dataset,
            validation_data=val_dataset,
            epochs=epochs,
            callbacks=callbacks_list
        )

        logging.info("Training completed.")
        model.save("checkpoints/final_model.keras")


if __name__ == "__main__":
    main()

总结

本代码实现了 Wide & Deep 推荐模型,适用于隐式反馈下的推荐或广告点击率预测任务。你可以根据实际业务需求进一步优化模型结构、加入负采样、使用更复杂网络、迁移到真实数据集等。


DeepFM

简介:DeepFM 推荐模型

  • 提出者:哈工大 & 4Paradigm
  • 年份:2017
  • 核心思想
    • Factorization Machine (FM)Deep Neural Network (DNN) 结合
    • FM 部分建模特征之间的二阶组合关系(显式)
    • Deep 部分建模高阶非线性组合(隐式)
    • 共享底层 Embedding 层,提高效率和泛化能力
  • 输入结构
    • dense_inputs:稠密特征(如用户年龄、物品价格等)
    • sparse_inputs:稀疏类别特征(如用户 ID、物品 ID、标签等)
  • 典型应用场景
    • 广告点击率预估(CTR)
    • 推荐系统(协同过滤 + 上下文信息)
    • 搜索排序优化

模型结构简述

组件描述
Embedding 层所有稀疏特征映射为低维向量
FM Layer二阶交互项(所有嵌入向量的内积之和)
Deep Layer多层感知机(MLP),输入为拼接的嵌入向量
Output Layer将 Wide(FM)和 Deep 的输出相加并通过 sigmoid 得到预测概率
损失函数Binary Crossentropy(隐式反馈)

使用 TensorFlow/Keras 模板完成 DeepFM 实现

以下是基于你给的模板构建的 DeepFM 推荐模型:

import tensorflow as tf
from tensorflow.keras import layers, models, optimizers, callbacks, losses, metrics
import os
import logging
import time
import numpy as np


# ================== Step 1: 数据准备 ==================
def get_dataset(num_samples=10000, num_dense_features=5, num_sparse_features=3, vocab_size=1000, test_size=0.2):
    """
    构造一个模拟的 CTR/推荐数据集(稠密+稀疏特征)
    :param num_samples: 总样本数
    :param num_dense_features: 稠密特征数量
    :param num_sparse_features: 稀疏特征数量
    :param vocab_size: 稀疏特征词汇表大小
    :param test_size: 测试集比例
    """
    # 稠密特征:随机浮点数
    dense_data = np.random.rand(num_samples, num_dense_features).astype(np.float32)

    # 稀疏特征:随机整数表示类别
    sparse_data = np.random.randint(0, vocab_size, size=(num_samples, num_sparse_features))

    # 标签:模拟点击率(0或1)
    labels = np.random.randint(0, 2, size=num_samples)

    indices = np.random.permutation(num_samples)
    split_idx = int((1 - test_size) * num_samples)
    train_idx, val_idx = indices[:split_idx], indices[split_idx:]

    train_dense, train_sparse, train_labels = dense_data[train_idx], sparse_data[train_idx], labels[train_idx]
    val_dense, val_sparse, val_labels = dense_data[val_idx], sparse_data[val_idx], labels[val_idx]

    return ((train_dense, train_sparse), train_labels), ((val_dense, val_sparse), val_labels)


# ================== Step 2: 构建模型 ==================
class DeepFMLayer(layers.Layer):
    def __init__(self, embedding_dim=10, sparse_feature_num=3, hidden_units=[64, 32], **kwargs):
        super(DeepFMLayer, self).__init__(**kwargs)
        self.embedding_dim = embedding_dim
        self.sparse_feature_num = sparse_feature_num

        # Embedding Layer for sparse features
        self.embedding_layer = layers.Embedding(input_dim=10000, output_dim=embedding_dim)

        # FM Layer: First and Second order terms
        self.fm_first_dense = layers.Dense(1)
        self.fm_second_dense = layers.Dense(1)

        # Deep Layer
        self.deep_layers = [
            layers.Dense(units, activation='relu') for units in hidden_units
        ]
        self.final_dense = layers.Dense(1, activation='sigmoid')

    def call(self, inputs):
        dense_inputs, sparse_inputs = inputs

        # 1. FM Part
        sparse_embeddings = self.embedding_layer(sparse_inputs)  # shape: (batch_size, sparse_feat_num, embedding_dim)

        # First Order Term
        first_order_input = tf.reshape(sparse_embeddings, (-1, self.sparse_feature_num * self.embedding_dim))
        first_order = self.fm_first_dense(first_order_input)

        # Second Order Term
        sum_of_vectors = tf.reduce_sum(sparse_embeddings, axis=1)  # shape: (batch_size, embedding_dim)
        square_of_sums = tf.square(sum_of_vectors)  # shape: (batch_size, embedding_dim)
        sum_of_squares = tf.reduce_sum(tf.square(sparse_embeddings), axis=1)  # shape: (batch_size, embedding_dim)
        second_order = 0.5 * tf.reduce_sum(square_of_sums - sum_of_squares, axis=1, keepdims=True)

        fm_logits = first_order + second_order  # shape: (batch_size, 1)

        # 2. Deep Part
        deep_input = tf.concat([
            dense_inputs,
            tf.reshape(sparse_embeddings, (-1, self.sparse_feature_num * self.embedding_dim))
        ], axis=1)

        deep_logits = deep_input
        for layer in self.deep_layers:
            deep_logits = layer(deep_logits)

        # Final Output
        logits = self.final_dense(tf.concat([fm_logits, deep_logits], axis=1))
        return logits


class DeepFMModel(models.Model):
    def __init__(self, embedding_dim=10, sparse_feature_num=3, hidden_units=[64, 32], **kwargs):
        super(DeepFMModel, self).__init__(**kwargs)
        self.deepfm_layer = DeepFMLayer(embedding_dim, sparse_feature_num, hidden_units)

    def call(self, inputs):
        return self.deepfm_layer(inputs)


# ================== Step 3: 工具函数与配置 ==================
def setup_logger(log_dir="logs"):
    os.makedirs(log_dir, exist_ok=True)
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(levelname)s - %(message)s',
        handlers=[
            logging.FileHandler(f"{log_dir}/train_{time.strftime('%Y%m%d_%H%M%S')}.log"),
            logging.StreamHandler()
        ]
    )


# ================== Step 4: 主程序入口 ==================
def main():
    setup_logger()
    strategy = tf.distribute.MirroredStrategy()
    logging.info(f'Number of devices: {strategy.num_replicas_in_sync}')

    with strategy.scope():
        num_dense_features = 5
        num_sparse_features = 3
        vocab_size = 1000
        batch_size = 512
        epochs = 50
        embedding_dim = 10
        hidden_units = [64, 32]

        # 获取数据集(替换为你自己的 Dataset)
        ((train_dense, train_sparse), train_labels), ((val_dense, val_sparse), val_labels) = get_dataset(
            num_dense_features=num_dense_features,
            num_sparse_features=num_sparse_features,
            vocab_size=vocab_size
        )

        train_dataset = tf.data.Dataset.from_tensor_slices(((train_dense, train_sparse), train_labels)) \
            .shuffle(1000).batch(batch_size).prefetch(tf.data.AUTOTUNE)
        val_dataset = tf.data.Dataset.from_tensor_slices(((val_dense, val_sparse), val_labels)) \
            .batch(batch_size).prefetch(tf.data.AUTOTUNE)

        # 构建模型
        model = DeepFMModel(embedding_dim=embedding_dim, sparse_feature_num=num_sparse_features, hidden_units=hidden_units)

        # 编译模型
        optimizer = optimizers.Adam(learning_rate=1e-3)
        loss_fn = losses.BinaryCrossentropy()
        metric = metrics.BinaryAccuracy(name='acc')

        model.compile(optimizer=optimizer, loss=loss_fn, metrics=[metric])

        # 回调函数
        checkpoint_cb = callbacks.ModelCheckpoint(
            "checkpoints/best_model.keras",
            monitor='val_loss',
            save_best_only=True,
            mode='min'
        )
        early_stop_cb = callbacks.EarlyStopping(patience=10, restore_best_weights=True)
        tensorboard_cb = callbacks.TensorBoard(log_dir="runs")

        callbacks_list = [checkpoint_cb, early_stop_cb, tensorboard_cb]

        # 开始训练
        history = model.fit(
            train_dataset,
            validation_data=val_dataset,
            epochs=epochs,
            callbacks=callbacks_list
        )

        logging.info("Training completed.")
        model.save("checkpoints/final_model.keras")


if __name__ == "__main__":
    main()

总结

本代码实现了 DeepFM 推荐模型,适用于广告点击率预估(CTR)和推荐系统任务。你可以根据实际业务需求进一步优化模型结构、加入负采样、使用更复杂网络、迁移到真实数据集等。


YouTube DNN Recommender

简介:YouTube DNN 推荐模型

  • 提出者:YouTube 团队(Google Research)
  • 年份:2016
  • 核心思想
    • 使用深度神经网络对用户行为建模,用于视频推荐
    • 模型分为两个阶段:
      • 召回阶段(Candidate Generation):从海量视频中快速筛选候选集
      • 排序阶段(Ranking):对候选集进行精排打分
    • 本实现主要关注 排序模型(Ranking Model)
  • 输入特征
    • 用户历史观看的视频 Embedding(平均池化)
    • 用户其他特征(如地区、性别等)
    • 当前候选视频的特征(如类别、标签、Embedding)
  • 输出目标
    • 对每个候选视频预测点击概率(CTR),用于排序

模型结构简述

组件描述
用户历史视频嵌入多个视频 ID 转为 Embedding 后平均池化
用户上下文特征如地区、设备、时间等稀疏/稠密特征
候选视频特征视频 ID、标签、时长等特征
Deep LayerMLP 层堆叠,融合所有特征
输出层Sigmoid 得到点击概率
损失函数Binary Crossentropy(隐式反馈)

使用 TensorFlow/Keras 模板完成 YouTube DNN 实现

以下是基于你给的模板构建的 YouTube DNN Ranking 推荐模型:

import tensorflow as tf
from tensorflow.keras import layers, models, optimizers, callbacks, losses, metrics
import os
import logging
import time
import numpy as np


# ================== Step 1: 数据准备 ==================
def get_dataset(num_samples=10000, max_history_length=10, num_videos=5000, test_size=0.2):
    """
    构造一个模拟的推荐数据集(用户历史 + 候选视频)
    :param num_samples: 总样本数
    :param max_history_length: 用户历史最大长度
    :param num_videos: 视频总数
    :param test_size: 测试集比例
    """
    # 用户历史观看记录(最多 max_history_length 个视频)
    user_histories = np.random.randint(0, num_videos, size=(num_samples, max_history_length))
    history_lengths = np.random.randint(1, max_history_length + 1, size=num_samples)

    # 候选视频 ID
    candidate_ids = np.random.randint(0, num_videos, size=num_samples)

    # 用户稠密特征(可扩展)
    user_dense_features = np.random.rand(num_samples, 3).astype(np.float32)

    # 标签:是否点击(二分类)
    labels = np.random.randint(0, 2, size=num_samples)

    indices = np.random.permutation(num_samples)
    split_idx = int((1 - test_size) * num_samples)
    train_idx, val_idx = indices[:split_idx], indices[split_idx:]

    train_data = (
        user_histories[train_idx],
        user_dense_features[train_idx],
        candidate_ids[train_idx]
    )
    train_labels = labels[train_idx]

    val_data = (
        user_histories[val_idx],
        user_dense_features[val_idx],
        candidate_ids[val_idx]
    )
    val_labels = labels[val_idx]

    return (train_data, train_labels), (val_data, val_labels), num_videos


# ================== Step 2: 构建模型 ==================
class YouTubeDNNModel(models.Model):
    def __init__(self, num_videos, embedding_dim=32, hidden_units=[128, 64], max_history_length=10, **kwargs):
        super(YouTubeDNNModel, self).__init__(**kwargs)
        self.embedding_dim = embedding_dim
        self.max_history_length = max_history_length

        # 视频 Embedding 层
        self.video_embedding = layers.Embedding(
            input_dim=num_videos,
            output_dim=embedding_dim,
            name='video_embedding'
        )

        # 用户历史行为处理(平均池化)
        self.user_history_pooling = layers.GlobalAveragePooling1D()

        # 用户稠密特征处理
        self.user_dense_proj = layers.Dense(embedding_dim)

        # 候选视频 Embedding
        self.candidate_proj = layers.Dense(embedding_dim)

        # Deep 部分
        self.deep_layers = [
            layers.Dense(units, activation='relu') for units in hidden_units
        ]
        self.final_dense = layers.Dense(1, activation='sigmoid')

    def call(self, inputs):
        histories, user_dense, candidates = inputs

        # 用户历史视频 Embedding 平均池化
        hist_embeddings = self.video_embedding(histories)
        user_hist_emb = self.user_history_pooling(hist_embeddings)

        # 用户稠密特征转换
        user_dense_emb = self.user_dense_proj(user_dense)

        # 候选视频 Embedding
        candidate_emb = self.video_embedding(candidates)

        # 特征拼接
        combined = tf.concat([
            user_hist_emb,
            user_dense_emb,
            candidate_emb
        ], axis=-1)

        # Deep 部分
        x = combined
        for layer in self.deep_layers:
            x = layer(x)

        # 输出点击率预测
        logits = self.final_dense(x)
        return logits


# ================== Step 3: 工具函数与配置 ==================
def setup_logger(log_dir="logs"):
    os.makedirs(log_dir, exist_ok=True)
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(levelname)s - %(message)s',
        handlers=[
            logging.FileHandler(f"{log_dir}/train_{time.strftime('%Y%m%d_%H%M%S')}.log"),
            logging.StreamHandler()
        ]
    )


# ================== Step 4: 主程序入口 ==================
def main():
    setup_logger()
    strategy = tf.distribute.MirroredStrategy()
    logging.info(f'Number of devices: {strategy.num_replicas_in_sync}')

    with strategy.scope():
        max_history_length = 10
        num_videos = 5000
        batch_size = 512
        epochs = 50
        embedding_dim = 32
        hidden_units = [128, 64]

        # 获取数据集(替换为你自己的 Dataset)
        ((train_hist, train_dense, train_candidate), train_labels), \
        ((val_hist, val_dense, val_candidate), val_labels), num_videos = get_dataset(
            max_history_length=max_history_length,
            num_videos=num_videos
        )

        train_dataset = tf.data.Dataset.from_tensor_slices((
            (train_hist, train_dense, train_candidate),
            train_labels
        )).shuffle(1000).batch(batch_size).prefetch(tf.data.AUTOTUNE)

        val_dataset = tf.data.Dataset.from_tensor_slices((
            (val_hist, val_dense, val_candidate),
            val_labels
        )).batch(batch_size).prefetch(tf.data.AUTOTUNE)

        # 构建模型
        model = YouTubeDNNModel(
            num_videos=num_videos,
            embedding_dim=embedding_dim,
            hidden_units=hidden_units,
            max_history_length=max_history_length
        )

        # 编译模型
        optimizer = optimizers.Adam(learning_rate=1e-3)
        loss_fn = losses.BinaryCrossentropy()
        metric = metrics.BinaryAccuracy(name='acc')

        model.compile(optimizer=optimizer, loss=loss_fn, metrics=[metric])

        # 回调函数
        checkpoint_cb = callbacks.ModelCheckpoint(
            "checkpoints/best_model.keras",
            monitor='val_loss',
            save_best_only=True,
            mode='min'
        )
        early_stop_cb = callbacks.EarlyStopping(patience=10, restore_best_weights=True)
        tensorboard_cb = callbacks.TensorBoard(log_dir="runs")

        callbacks_list = [checkpoint_cb, early_stop_cb, tensorboard_cb]

        # 开始训练
        history = model.fit(
            train_dataset,
            validation_data=val_dataset,
            epochs=epochs,
            callbacks=callbacks_list
        )

        logging.info("Training completed.")
        model.save("checkpoints/final_model.keras")


if __name__ == "__main__":
    main()

总结

本代码实现了 YouTube DNN Ranking 推荐模型,适用于广告点击率预估(CTR)、视频推荐排序任务。你可以根据实际业务需求进一步优化模型结构、加入负采样、使用更复杂网络、迁移到真实数据集等。

Graph Neural Networks (GNNs) for Recommendation

简介:GNN for Recommendation

  • 提出背景:
    • 用户与物品之间、物品与物品之间的关系天然适合用图结构建模
    • GNN 可以通过消息传递机制聚合邻居节点信息,从而学习更丰富的嵌入表示
  • 典型模型:
    • NGCF (Neural Graph Collaborative Filtering):基于 GNN 的协同过滤
    • LightGCN:简化版 GCN,去除非线性激活和特征变换,仅保留邻域传播
    • PinSage / DGRec / DGT:引入采样机制处理大规模图数据
  • 输入结构:
    • 图结构:用户 ↔ 物品的交互边 + 用户/物品属性(可选)
  • 模型输出:用户和物品的嵌入向量,用于 Top-K 推荐或 CTR 预测

模型结构简述(以 LightGCN 为例)

组件描述
图构建构建用户-物品二部图
多层 GNN使用图卷积层传播信息(无参数)
嵌入生成最终嵌入为各层嵌入的加权平均
推荐预测用户嵌入与物品嵌入点积作为评分

使用 TensorFlow/Keras 模板完成 GNN-based 推荐模型(以 LightGCN 为例)

以下是一个完整的 LightGCN 推荐模型实现,支持训练、验证、保存、可视化等功能:

import tensorflow as tf
from tensorflow.keras import layers, models, optimizers, callbacks, losses, metrics
import numpy as np
import os
import logging
import time
import networkx as nx
from collections import defaultdict


# ================== Step 1: 数据准备 ==================
def build_interaction_graph(user_item_pairs, num_users, num_items):
    """
    构造用户-物品二部图的邻接矩阵
    :param user_item_pairs: List of (user_id, item_id)
    :return: adjacency matrix (num_nodes x num_nodes)
    """
    graph = nx.Graph()

    # 添加所有节点
    for u in range(num_users):
        graph.add_node(u, type='user')
    for i in range(num_users, num_users + num_items):
        graph.add_node(i, type='item')

    # 添加边
    for u, i in user_item_pairs:
        graph.add_edge(u, num_users + i)

    adj_matrix = nx.adjacency_matrix(graph).toarray()
    return tf.constant(adj_matrix, dtype=tf.float32)


def get_dataset(num_users=1000, num_items=500, interaction_density=0.01, test_size=0.2):
    """
    构造用户-物品交互数据集
    """
    interactions = []
    for u in range(num_users):
        num_interactions = int(interaction_density * num_items)
        items = np.random.choice(num_items, size=num_interactions, replace=False)
        interactions.extend([(u, i) for i in items])

    labels = np.ones(len(interactions))
    indices = np.random.permutation(len(interactions))
    split_idx = int((1 - test_size) * len(interactions))
    train_idx, val_idx = indices[:split_idx], indices[split_idx:]

    train_pairs = [interactions[i] for i in train_idx]
    val_pairs = [interactions[i] for i in val_idx]

    # 构造邻接矩阵
    A = build_interaction_graph(train_pairs, num_users, num_items)

    return train_pairs, val_pairs, A, num_users, num_items


# ================== Step 2: 构建模型 ==================
class LightGCNLayer(layers.Layer):
    def __init__(self, A, embedding_dim=64, num_layers=3, **kwargs):
        super(LightGCNLayer, self).__init__(**kwargs)
        self.A = A
        self.num_layers = num_layers
        self.embedding_dim = embedding_dim

        self.user_embeddings = tf.keras.initializers.RandomNormal()((num_users + num_items, embedding_dim))

    def call(self, inputs=None):
        embeddings = self.user_embeddings
        all_embeddings = [embeddings]

        for _ in range(self.num_layers):
            embeddings = tf.matmul(self.A, embeddings)
            all_embeddings.append(embeddings)

        # 多层嵌入平均
        final_embeddings = tf.reduce_mean(tf.stack(all_embeddings, axis=0), axis=0)
        return final_embeddings


class GNNRecommender(models.Model):
    def __init__(self, A, num_users, num_items, embedding_dim=64, num_layers=3, **kwargs):
        super(GNNRecommender, self).__init__(**kwargs)
        self.gnn_layer = LightGCNLayer(A=A, embedding_dim=embedding_dim, num_layers=num_layers)
        self.num_users = num_users
        self.num_items = num_items

    def call(self, inputs):
        user_ids, item_ids = inputs
        embeddings = self.gnn_layer()
        user_emb = tf.nn.embedding_lookup(embeddings, user_ids)
        item_emb = tf.nn.embedding_lookup(embeddings, self.num_users + item_ids)
        logits = tf.reduce_sum(user_emb * item_emb, axis=1, keepdims=True)
        return logits

    def get_top_k(self, user_ids, k=10):
        embeddings = self.gnn_layer()
        user_embs = tf.nn.embedding_lookup(embeddings, user_ids)
        item_embs = embeddings[self.num_users:]
        scores = tf.matmul(user_embs, item_embs, transpose_b=True)
        top_k_indices = tf.math.top_k(scores, k=k).indices
        return top_k_indices


# ================== Step 3: 工具函数与配置 ==================
def setup_logger(log_dir="logs"):
    os.makedirs(log_dir, exist_ok=True)
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(levelname)s - %(message)s',
        handlers=[
            logging.FileHandler(f"{log_dir}/train_{time.strftime('%Y%m%d_%H%M%S')}.log"),
            logging.StreamHandler()
        ]
    )


# ================== Step 4: 主程序入口 ==================
def main():
    setup_logger()
    strategy = tf.distribute.MirroredStrategy()
    logging.info(f'Number of devices: {strategy.num_replicas_in_sync}')

    with strategy.scope():
        num_users = 1000
        num_items = 500
        batch_size = 1024
        epochs = 50
        embedding_dim = 64
        num_layers = 3

        # 获取数据集
        train_pairs, val_pairs, A, num_users, num_items = get_dataset(num_users=num_users, num_items=num_items)

        # 构造训练数据
        train_users = np.array([p[0] for p in train_pairs])
        train_items = np.array([p[1] for p in train_pairs])
        train_labels = np.ones(len(train_pairs))

        val_users = np.array([p[0] for p in val_pairs])
        val_items = np.array([p[1] for p in val_pairs])
        val_labels = np.ones(len(val_pairs))

        train_dataset = tf.data.Dataset.from_tensor_slices(((train_users, train_items), train_labels)) \
            .shuffle(10000).batch(batch_size).prefetch(tf.data.AUTOTUNE)
        val_dataset = tf.data.Dataset.from_tensor_slices(((val_users, val_items), val_labels)) \
            .batch(batch_size).prefetch(tf.data.AUTOTUNE)

        # 构建模型
        model = GNNRecommender(A=A, num_users=num_users, num_items=num_items, embedding_dim=embedding_dim, num_layers=num_layers)

        # 编译模型
        optimizer = optimizers.Adam(learning_rate=1e-3)
        loss_fn = losses.BinaryCrossentropy(from_logits=True)
        metric = metrics.BinaryAccuracy(name='acc')

        model.compile(optimizer=optimizer, loss=loss_fn, metrics=[metric])

        # 回调函数
        checkpoint_cb = callbacks.ModelCheckpoint(
            "checkpoints/best_model.keras",
            monitor='val_loss',
            save_best_only=True,
            mode='min'
        )
        early_stop_cb = callbacks.EarlyStopping(patience=10, restore_best_weights=True)
        tensorboard_cb = callbacks.TensorBoard(log_dir="runs")

        callbacks_list = [checkpoint_cb, early_stop_cb, tensorboard_cb]

        # 开始训练
        history = model.fit(
            train_dataset,
            validation_data=val_dataset,
            epochs=epochs,
            callbacks=callbacks_list
        )

        logging.info("Training completed.")
        model.save("checkpoints/final_model.keras")


if __name__ == "__main__":
    main()

总结

本代码实现了 基于图神经网络(LightGCN)的推荐系统模型,适用于社交推荐、商品推荐、视频推荐等任务。你可以根据实际业务需求进一步优化模型结构、加入负采样、使用更复杂网络、迁移到真实数据集等。

Recurrent Neural Networks (RNNs) for Session-based Recommendation

简介:Session-based 推荐系统

  • 核心任务:
    • 给定用户在当前会话中点击/浏览的一系列物品(item sequence),预测下一个最可能点击的物品。
    • 没有长期用户 ID 的信息,只依赖短期行为序列建模。
  • 典型模型:
    • GRU4Rec(2016):最早将 GRU 应用于会话推荐
    • NARM(2017):结合 GRU + Attention,提取兴趣点
    • STAMP(2018):不使用 RNN,用 attention 提取短期兴趣
    • SASRec(2019):Transformer 自注意力机制建模序列
  • 输入结构:
    • session_items:每个会话中的物品 ID 序列(变长)
  • 输出:下一个物品的概率分布(Top-K 推荐)
  • 损失函数:
    • 使用交叉熵(CrossEntropy)预测下一个 item
    • 或者负采样 + 点积打分(类似 GNN)

模型结构简述(以 GRU4Rec 为例)

组件描述
Embedding Layer将物品 ID 映射为低维向量
RNN / GRU Layer编码会话行为序列
Output Layer预测下一个 item 的概率分布
Loss FunctionCrossEntropyLoss 或 SampledSoftmax

使用 TensorFlow/Keras 模板完成 Session-based RNN 推荐模型(以 GRU4Rec 为例)

以下是一个完整的 GRU4Rec 模型实现,支持训练、验证、保存、可视化等功能:

import tensorflow as tf
from tensorflow.keras import layers, models, optimizers, callbacks, losses, metrics
import numpy as np
import os
import logging
import time


# ================== Step 1: 数据准备 ==================
def generate_session_data(num_sessions=10000, max_seq_length=20, num_items=5000, test_size=0.2):
    """
    构造模拟的会话数据集(Session-based)
    :param num_sessions: 总会话数
    :param max_seq_length: 最大序列长度
    :param num_items: 物品总数
    :param test_size: 测试集比例
    """
    session_sequences = []
    next_item_labels = []

    for _ in range(num_sessions):
        seq_len = np.random.randint(3, max_seq_length + 1)
        seq = np.random.randint(0, num_items, size=seq_len)
        session_sequences.append(seq[:-1])  # 前 n-1 items
        next_item_labels.append(seq[-1])   # 第 n item 是目标

    # Padding
    session_sequences = tf.keras.preprocessing.sequence.pad_sequences(
        session_sequences,
        maxlen=max_seq_length - 1,
        padding='post',
        truncating='post'
    )
    next_item_labels = np.array(next_item_labels)

    indices = np.random.permutation(num_sessions)
    split_idx = int((1 - test_size) * num_sessions)
    train_idx, val_idx = indices[:split_idx], indices[split_idx:]

    X_train, y_train = session_sequences[train_idx], next_item_labels[train_idx]
    X_val, y_val = session_sequences[val_idx], next_item_labels[val_idx]

    return (X_train, y_train), (X_val, y_val), num_items


# ================== Step 2: 构建模型 ==================
class RNN4RecModel(models.Model):
    def __init__(self, num_items, embedding_dim=64, hidden_units=128, max_seq_length=20, **kwargs):
        super(RNN4RecModel, self).__init__(**kwargs)
        self.embedding_layer = layers.Embedding(input_dim=num_items, output_dim=embedding_dim)
        self.rnn_layer = layers.GRU(hidden_units, return_sequences=False)
        self.output_proj = layers.Dense(num_items)

    def call(self, inputs):
        x = self.embedding_layer(inputs)
        x = self.rnn_layer(x)
        logits = self.output_proj(x)
        return logits

    def get_top_k(self, session_seq, k=10):
        """
        输入一个 session 序列,输出 top-k 推荐
        """
        padded_seq = tf.keras.preprocessing.sequence.pad_sequences(
            [session_seq], maxlen=self.input_shape[1], padding='post')
        logits = self.predict(padded_seq, verbose=0)
        top_k_indices = tf.math.top_k(logits, k=k).indices.numpy()[0]
        return top_k_indices


# ================== Step 3: 工具函数与配置 ==================
def setup_logger(log_dir="logs"):
    os.makedirs(log_dir, exist_ok=True)
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(levelname)s - %(message)s',
        handlers=[
            logging.FileHandler(f"{log_dir}/train_{time.strftime('%Y%m%d_%H%M%S')}.log"),
            logging.StreamHandler()
        ]
    )


# ================== Step 4: 主程序入口 ==================
def main():
    setup_logger()
    strategy = tf.distribute.MirroredStrategy()
    logging.info(f'Number of devices: {strategy.num_replicas_in_sync}')

    with strategy.scope():
        num_items = 5000
        batch_size = 512
        epochs = 50
        embedding_dim = 64
        hidden_units = 128
        max_seq_length = 20

        # 获取数据集
        (X_train, y_train), (X_val, y_val), num_items = generate_session_data(
            num_sessions=20000,
            max_seq_length=max_seq_length,
            num_items=num_items
        )

        train_dataset = tf.data.Dataset.from_tensor_slices((X_train, y_train)) \
            .shuffle(10000).batch(batch_size).prefetch(tf.data.AUTOTUNE)
        val_dataset = tf.data.Dataset.from_tensor_slices((X_val, y_val)) \
            .batch(batch_size).prefetch(tf.data.AUTOTUNE)

        # 构建模型
        model = RNN4RecModel(num_items=num_items, embedding_dim=embedding_dim,
                             hidden_units=hidden_units, max_seq_length=max_seq_length)

        # 编译模型
        optimizer = optimizers.Adam(learning_rate=1e-3)
        loss_fn = losses.SparseCategoricalCrossentropy(from_logits=True)
        metric = metrics.SparseCategoricalAccuracy(name='acc')

        model.compile(optimizer=optimizer, loss=loss_fn, metrics=[metric])

        # 回调函数
        checkpoint_cb = callbacks.ModelCheckpoint(
            "checkpoints/best_model.keras",
            monitor='val_loss',
            save_best_only=True,
            mode='min'
        )
        early_stop_cb = callbacks.EarlyStopping(patience=10, restore_best_weights=True)
        tensorboard_cb = callbacks.TensorBoard(log_dir="runs")

        callbacks_list = [checkpoint_cb, early_stop_cb, tensorboard_cb]

        # 开始训练
        history = model.fit(
            train_dataset,
            validation_data=val_dataset,
            epochs=epochs,
            callbacks=callbacks_list
        )

        logging.info("Training completed.")
        model.save("checkpoints/final_model.keras")


if __name__ == "__main__":
    main()

总结

本代码实现了 基于 RNN(GRU)的会话推荐系统模型(Session-based Recommender),适用于电商、视频平台、新闻资讯等场景下的无用户 ID 的推荐任务。你可以根据实际业务需求进一步优化模型结构、加入 Attention、使用 Transformer、迁移到真实数据集等。

Transformer-based Models (BERT4Rec)

简介:Transformer-based 推荐系统

  • 核心思想:
    • 使用 Transformer 构建用户行为序列的双向上下文感知表示
    • 在训练时采用”掩码语言模型”(Masked Language Model, MLM)任务,预测被掩码的行为项(item)
    • 推理时使用最后一层表示预测下一个 item 或 Top-K 推荐
  • 典型模型:
    • BERT4Rec(2019):首个将 BERT 思想用于推荐系统的模型
    • SASRec(Self-Attentive Sequential Recommendation)
    • DuoRec / UniSRec / M6Rec:多任务、预训练方向发展模型
  • 输入结构:
    • input_ids:物品 ID 序列(如 [i1, i2, i3, i4]
    • attention_mask:指示哪些位置是填充的
  • 输出:每个位置对应的 item 表示,用于 MLM 预测或 Top-K 推荐

模型结构简述(BERT4Rec)

组件描述
Embedding Layer包括 token 和 position embedding
Transformer Encoder多层自注意力模块编码序列
MLM Head预测被 mask 的 item
Loss FunctionCrossEntropyLoss on masked positions

使用 TensorFlow/Keras 模板完成 BERT4Rec 推荐模型

以下是一个完整的 BERT4Rec 推荐模型实现,支持训练、验证、保存、可视化等功能:

import tensorflow as tf
from tensorflow.keras import layers, models, optimizers, callbacks, losses, metrics
import numpy as np
import os
import logging
import time


# ================== Step 1: 数据准备 ==================
def generate_sequence_data(num_sequences=10000, max_seq_length=50, num_items=5000, test_size=0.2, mask_prob=0.15):
    """
    构造模拟的序列数据集(Sequence-based),并进行 MLM 掩码
    :param num_sequences: 总样本数
    :param max_seq_length: 最大序列长度
    :param num_items: 物品总数
    :param test_size: 测试集比例
    :param mask_prob: 掩码概率
    """
    sequences = []
    for _ in range(num_sequences):
        seq_len = np.random.randint(5, max_seq_length + 1)
        seq = np.random.randint(0, num_items, size=seq_len)
        sequences.append(seq)

    # Padding
    input_ids = tf.keras.preprocessing.sequence.pad_sequences(
        sequences,
        maxlen=max_seq_length,
        padding='post',
        truncating='post'
    )

    # 构造 MLM 输入和标签
    labels = input_ids.copy()
    attention_masks = (input_ids != 0).astype(np.int32)

    # 随机掩码
    mask = np.random.rand(*input_ids.shape) < mask_prob
    mask &= (input_ids != 0)  # 不对 padding 进行掩码
    input_ids[mask] = 0  # 用 0 表示 [MASK]

    indices = np.random.permutation(num_sequences)
    split_idx = int((1 - test_size) * num_sequences)
    train_idx, val_idx = indices[:split_idx], indices[split_idx:]

    return (
        (input_ids[train_idx], attention_masks[train_idx], labels[train_idx]),
        (input_ids[val_idx], attention_masks[val_idx], labels[val_idx]),
        num_items
    )


# ================== Step 2: 构建模型 ==================
class TransformerEmbedding(layers.Layer):
    def __init__(self, vocab_size, embedding_dim=128, max_position_embeddings=50, **kwargs):
        super(TransformerEmbedding, self).__init__(**kwargs)
        self.token_embedding = layers.Embedding(vocab_size, embedding_dim)
        self.position_embedding = layers.Embedding(max_position_embeddings, embedding_dim)
        self.layer_norm = layers.LayerNormalization(epsilon=1e-12)
        self.dropout = layers.Dropout(0.1)

    def call(self, inputs, training=False):
        input_ids, attention_mask = inputs
        seq_length = tf.shape(input_ids)[1]
        position_ids = tf.range(seq_length, dtype=tf.int32)[tf.newaxis, :]
        token_emb = self.token_embedding(input_ids)
        pos_emb = self.position_embedding(position_ids)
        embeddings = token_emb + pos_emb
        embeddings = self.layer_norm(embeddings)
        embeddings = self.dropout(embeddings, training=training)
        return embeddings


class BERT4RecModel(models.Model):
    def __init__(self, vocab_size, embedding_dim=128, num_layers=4, num_heads=8, ff_dim=256, max_seq_length=50, **kwargs):
        super(BERT4RecModel, self).__init__(**kwargs)
        self.embedding_layer = TransformerEmbedding(vocab_size, embedding_dim, max_seq_length)
        self.transformer_blocks = [
            tf.keras.Sequential([
                layers.MultiHeadAttention(num_heads=num_heads, key_dim=embedding_dim),
                layers.LayerNormalization(),
                layers.Dense(ff_dim, activation='relu'),
                layers.Dense(embedding_dim),
                layers.LayerNormalization()
            ]) for _ in range(num_layers)
        ]
        self.mlm_head = layers.Dense(vocab_size)

    def call(self, inputs, training=False):
        input_ids, attention_mask = inputs
        x = self.embedding_layer((input_ids, attention_mask), training=training)

        # Attention Mask
        attention_weights = tf.cast(tf.expand_dims(attention_mask, axis=1), tf.float32)
        attention_weights = tf.matmul(attention_weights, attention_weights, transpose_b=True)

        for block in self.transformer_blocks:
            x = block([x, x, x, attention_weights], training=training)

        mlm_logits = self.mlm_head(x)
        return mlm_logits

    def get_top_k(self, session_seq, k=10):
        padded_seq = tf.keras.preprocessing.sequence.pad_sequences(
            [session_seq], maxlen=self.input_shape[1][1], padding='post')
        logits = self.predict(padded_seq, verbose=0)
        top_k_indices = tf.math.top_k(logits[0][-1], k=k).indices.numpy()
        return top_k_indices


# ================== Step 3: 工具函数与配置 ==================
def setup_logger(log_dir="logs"):
    os.makedirs(log_dir, exist_ok=True)
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(levelname)s - %(message)s',
        handlers=[
            logging.FileHandler(f"{log_dir}/train_{time.strftime('%Y%m%d_%H%M%S')}.log"),
            logging.StreamHandler()
        ]
    )


# ================== Step 4: 主程序入口 ==================
def main():
    setup_logger()
    strategy = tf.distribute.MirroredStrategy()
    logging.info(f'Number of devices: {strategy.num_replicas_in_sync}')

    with strategy.scope():
        num_items = 5000
        batch_size = 256
        epochs = 50
        embedding_dim = 128
        num_layers = 4
        num_heads = 8
        ff_dim = 256
        max_seq_length = 50

        # 获取数据集
        (X_train_ids, X_train_mask, y_train), (X_val_ids, X_val_mask, y_val), num_items = generate_sequence_data(
            num_sequences=20000,
            max_seq_length=max_seq_length,
            num_items=num_items
        )

        train_dataset = tf.data.Dataset.from_tensor_slices(((X_train_ids, X_train_mask), y_train)) \
            .shuffle(10000).batch(batch_size).prefetch(tf.data.AUTOTUNE)
        val_dataset = tf.data.Dataset.from_tensor_slices(((X_val_ids, X_val_mask), y_val)) \
            .batch(batch_size).prefetch(tf.data.AUTOTUNE)

        # 构建模型
        model = BERT4RecModel(
            vocab_size=num_items,
            embedding_dim=embedding_dim,
            num_layers=num_layers,
            num_heads=num_heads,
            ff_dim=ff_dim,
            max_seq_length=max_seq_length
        )

        # 编译模型
        optimizer = optimizers.AdamW(learning_rate=1e-4, weight_decay=0.01)
        loss_fn = losses.SparseCategoricalCrossentropy(from_logits=True)
        metric = metrics.SparseCategoricalAccuracy(name='acc')

        model.compile(optimizer=optimizer, loss=loss_fn, metrics=[metric])

        # 回调函数
        checkpoint_cb = callbacks.ModelCheckpoint(
            "checkpoints/best_model.keras",
            monitor='val_loss',
            save_best_only=True,
            mode='min'
        )
        early_stop_cb = callbacks.EarlyStopping(patience=10, restore_best_weights=True)
        tensorboard_cb = callbacks.TensorBoard(log_dir="runs")

        callbacks_list = [checkpoint_cb, early_stop_cb, tensorboard_cb]

        # 开始训练
        history = model.fit(
            train_dataset,
            validation_data=val_dataset,
            epochs=epochs,
            callbacks=callbacks_list
        )

        logging.info("Training completed.")
        model.save("checkpoints/final_model.keras")


if __name__ == "__main__":
    main()

总结

本代码实现了 基于 Transformer 的推荐系统模型(BERT4Rec),适用于电商、视频平台、新闻资讯等场景下的序列推荐任务。你可以根据实际业务需求进一步优化模型结构、加入负采样、使用更复杂网络、迁移到真实数据集等。

Long Short-Term Memory(LSTM)

LSTM 是一种特殊的循环神经网络(RNN)结构,专门设计用来解决传统 RNN 在处理长序列时的梯度消失和梯度爆炸问题

它能够记住长期依赖信息,在自然语言处理、语音识别、时间序列预测等任务中表现优异。

传统的 RNN 在理论上可以记忆前面的信息,但在实践中,当序列变长时,模型很难记住早期的信息。这是因为反向传播过程中,梯度会逐渐变小(梯度消失),导致模型无法学习到长期依赖关系。

LSTM 的核心思想:门控机制

LSTM 引入了”门控机制(Gating Mechanism)“,通过三个主要的”门”来控制信息的流动:

  1. 遗忘门(Forget Gate):决定哪些信息从 cell state(记忆单元)中被丢弃。

    • 使用 sigmoid 激活函数输出一个 0~1 的值:
      • 0 表示完全忘记
      • 1 表示完全保留
  2. 输入门(Input Gate):决定哪些新信息将被加入到 cell state 中。

    • 包括两个部分:
      • 一个 sigmoid 层决定哪些值要更新
      • 一个 tanh 层创建一个新的候选值向量
  3. 输出门(Output Gate):决定当前时刻的输出是什么。

    • 首先用 sigmoid 层决定 cell state 的哪部分将被输出
    • 然后将 cell state 通过 tanh 激活函数,并与 sigmoid 输出相乘,得到最终输出

使用 TensorFlow 模板完成 LSTM 实现

import tensorflow as tf
from tensorflow.keras import layers, models, optimizers, callbacks, losses, metrics
import os
import logging
import time
import numpy as np


# ================== 自定义LSTM组件 ==================
class CustomLSTMCell(layers.Layer):
    def __init__(self, units, **kwargs):
        super(CustomLSTMCell, self).__init__(**kwargs)
        self.units = units
        self.state_size = [tf.TensorShape([units]), tf.TensorShape([units])]  # [h, c]
        self.output_size = units

    def build(self, input_shape):
        input_dim = input_shape[-1]

        # 输入门权重
        self.W_i = self.add_weight(
            shape=(input_dim, self.units),
            name='W_i',
            initializer='glorot_uniform')
        self.U_i = self.add_weight(
            shape=(self.units, self.units),
            name='U_i',
            initializer='orthogonal')
        self.b_i = self.add_weight(
            shape=(self.units,),
            name='b_i',
            initializer='zeros')

        # 遗忘门权重
        self.W_f = self.add_weight(
            shape=(input_dim, self.units),
            name='W_f',
            initializer='glorot_uniform')
        self.U_f = self.add_weight(
            shape=(self.units, self.units),
            name='U_f',
            initializer='orthogonal')
        self.b_f = self.add_weight(
            shape=(self.units,),
            name='b_f',
            initializer='ones')  # 初始偏置通常设为1

        # 候选记忆权重
        self.W_c = self.add_weight(
            shape=(input_dim, self.units),
            name='W_c',
            initializer='glorot_uniform')
        self.U_c = self.add_weight(
            shape=(self.units, self.units),
            name='U_c',
            initializer='orthogonal')
        self.b_c = self.add_weight(
            shape=(self.units,),
            name='b_c',
            initializer='zeros')

        # 输出门权重
        self.W_o = self.add_weight(
            shape=(input_dim, self.units),
            name='W_o',
            initializer='glorot_uniform')
        self.U_o = self.add_weight(
            shape=(self.units, self.units),
            name='U_o',
            initializer='orthogonal')
        self.b_o = self.add_weight(
            shape=(self.units,),
            name='b_o',
            initializer='zeros')

        self.built = True

    def call(self, inputs, states):
        h_prev, c_prev = states

        # 输入门
        i = tf.sigmoid(
            tf.matmul(inputs, self.W_i) +
            tf.matmul(h_prev, self.U_i) +
            self.b_i)

        # 遗忘门
        f = tf.sigmoid(
            tf.matmul(inputs, self.W_f) +
            tf.matmul(h_prev, self.U_f) +
            self.b_f)

        # 候选记忆
        c_hat = tf.tanh(
            tf.matmul(inputs, self.W_c) +
            tf.matmul(h_prev, self.U_c) +
            self.b_c)

        # 更新细胞状态
        c = f * c_prev + i * c_hat

        # 输出门
        o = tf.sigmoid(
            tf.matmul(inputs, self.W_o) +
            tf.matmul(h_prev, self.U_o) +
            self.b_o)

        # 新隐藏状态
        h = o * tf.tanh(c)

        return h, [h, c]

    def get_config(self):
        config = super(CustomLSTMCell, self).get_config()
        config.update({
            'units': self.units,
        })
        return config


class CustomLSTM(layers.Layer):
    def __init__(self, units, return_sequences=False, **kwargs):
        super(CustomLSTM, self).__init__(**kwargs)
        self.units = units
        self.return_sequences = return_sequences
        self.cell = CustomLSTMCell(units)

    def call(self, inputs):
        if isinstance(inputs, (list, tuple)):
            inputs, initial_state = inputs
        else:
            initial_state = None

        batch_size = tf.shape(inputs)[0]
        time_steps = tf.shape(inputs)[1]

        if initial_state is None:
            h = tf.zeros([batch_size, self.units], dtype=inputs.dtype)
            c = tf.zeros([batch_size, self.units], dtype=inputs.dtype)
            initial_state = [h, c]

        # 使用tf.TensorArray来收集所有时间步的输出
        outputs = tf.TensorArray(
            dtype=inputs.dtype,
            size=time_steps,
            dynamic_size=False)

        # 使用tf.while_loop代替Python循环
        def cond(t, *args):
            return t < time_steps

        def body(t, h, c, outputs):
            x_t = inputs[:, t, :]
            h, [h, c] = self.cell(x_t, [h, c])
            outputs = outputs.write(t, h)
            return t + 1, h, c, outputs

        _, h, c, outputs = tf.while_loop(
            cond=cond,
            body=body,
            loop_vars=(
                0,
                initial_state[0],
                initial_state[1],
                outputs)
        )

        # 将TensorArray转换为Tensor
        outputs = outputs.stack()  # [time, batch, units]
        outputs = tf.transpose(outputs, [1, 0, 2])  # [batch, time, units]

        if self.return_sequences:
            return outputs
        else:
            return outputs[:, -1, :]  # 只返回最后时间步的输出

    def get_config(self):
        config = super(CustomLSTM, self).get_config()
        config.update({
            'units': self.units,
            'return_sequences': self.return_sequences,
        })
        return config


# ================== 数据准备 ==================
def get_sequence_dataset(batch_size=32, seq_length=20, input_dim=10, num_classes=5):
    """生成模拟序列数据"""
    x_train = np.random.rand(1000, seq_length, input_dim).astype(np.float32)
    y_train = np.random.randint(0, num_classes, (1000,))
    x_val = np.random.rand(200, seq_length, input_dim).astype(np.float32)
    y_val = np.random.randint(0, num_classes, (200,))

    train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train))
    val_dataset = tf.data.Dataset.from_tensor_slices((x_val, y_val))

    train_dataset = train_dataset.shuffle(1000).batch(batch_size).prefetch(tf.data.AUTOTUNE)
    val_dataset = val_dataset.batch(batch_size).prefetch(tf.data.AUTOTUNE)

    return train_dataset, val_dataset


# ================== 构建模型 ==================
def build_lstm_model(seq_length=20, input_dim=10, num_classes=5):
    model = models.Sequential([
        layers.Input(shape=(seq_length, input_dim)),
        CustomLSTM(64, return_sequences=True),  # 自定义LSTM层
        CustomLSTM(32),
        layers.Dense(num_classes, activation='softmax')
    ])
    return model


# ================== Step 3: 工具函数与配置 ==================
def setup_logger(log_dir="logs"):
    os.makedirs(log_dir, exist_ok=True)
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(levelname)s - %(message)s',
        handlers=[
            logging.FileHandler(f"{log_dir}/train_{time.strftime('%Y%m%d_%H%M%S')}.log"),
            logging.StreamHandler()
        ]
    )


# ================== 主程序 ==================
def main():
    setup_logger()

    # 检查是否有GPU可用
    gpus = tf.config.list_physical_devices('GPU')
    if gpus:
        strategy = tf.distribute.MirroredStrategy()
    else:
        strategy = tf.distribute.get_strategy()  # 默认策略(CPU)

    logging.info(f'Number of devices: {strategy.num_replicas_in_sync}')

    with strategy.scope():
        batch_size = 64
        epochs = 10
        seq_length = 20
        input_dim = 10
        num_classes = 5

        # 获取数据集
        train_dataset, val_dataset = get_sequence_dataset(
            batch_size, seq_length, input_dim, num_classes)

        # 构建模型
        model = build_lstm_model(seq_length, input_dim, num_classes)

        # 编译模型
        model.compile(
            optimizer=optimizers.Adam(learning_rate=1e-3),
            loss=losses.SparseCategoricalCrossentropy(),
            metrics=[metrics.SparseCategoricalAccuracy(name='acc')]
        )

        # 回调函数
        callbacks_list = [
            callbacks.ModelCheckpoint(
                "checkpoints/best_lstm_model.keras",
                monitor='val_acc',
                save_best_only=True,
                mode='max'
            ),
            callbacks.EarlyStopping(patience=3, restore_best_weights=True),
            callbacks.TensorBoard(log_dir="runs/lstm")
        ]

        # 训练模型
        history = model.fit(
            train_dataset,
            validation_data=val_dataset,
            epochs=epochs,
            callbacks=callbacks_list
        )

        logging.info("Training completed.")
        model.save("checkpoints/final_lstm_model.keras")


if __name__ == "__main__":
    main()

总结

本代码实现了 自定义 LSTM 模型,从底层实现了 LSTM 的门控机制(遗忘门、输入门、输出门),使用 tf.while_loop 进行时间步展开。适用于序列分类、时间序列预测等任务。你可以根据实际业务需求进一步优化模型结构、调整层数和单元数、迁移到真实数据集等。

TensorFlow 所有常用层完整汇总表(tf.keras.layers)

一、基础层(Core Layers)

层名全称/说明作用输入形状输出形状示例
Dense全连接层线性变换:y=Wx+b(batch_size, ..., input_dim)(batch_size, ..., units)layers.Dense(64, activation='relu')
Activation激活函数应用非线性激活函数同上一层输出同上一层输出layers.Activation('relu')
Dropout随机失活训练时随机丢弃神经元同上一层输出同上一层输出layers.Dropout(0.5)
Input输入占位符定义模型输入格式根据任务而定同输入layers.Input(shape=(784,))

二、卷积层(Convolutional Layers)

层名全称作用输入形状输出形状示例
Conv1D一维卷积处理序列数据(batch_size, steps, channels)(batch_size, new_steps, filters)layers.Conv1D(32, 3, activation='relu')
Conv2D二维卷积图像特征提取(batch_size, height, width, channels)(batch_size, new_height, new_width, filters)layers.Conv2D(64, (3,3), activation='relu')
Conv3D三维卷积视频、医学图像处理(batch_size, depth, height, width, channels)(batch_size, new_depth, new_height, new_width, filters)layers.Conv3D(64, (3,3,3), activation='relu')
SeparableConv1D/2D可分离卷积减少计算量同 Conv1D/2D同 Conv1D/2Dlayers.SeparableConv2D(32, (3,3), activation='relu')
DepthwiseConv2D深度可分离卷积单通道卷积同 Conv2D同 Conv2Dlayers.DepthwiseConv2D(kernel_size=3)

三、池化层(Pooling Layers)

层名全称作用输入形状输出形状示例
MaxPooling1D/2D/3D最大池化下采样,保留最大值同对应卷积层减小尺寸后的相同格式layers.MaxPooling2D(pool_size=(2,2))
AveragePooling1D/2D/3D平均池化下采样,保留平均值同上同上layers.AveragePooling2D(pool_size=(2,2))
GlobalMaxPooling1D/2D/3D全局最大池化对整个特征图取最大值同上(batch_size, features)layers.GlobalMaxPooling2D()
GlobalAveragePooling1D/2D/3D全局平均池化对整个特征图取平均值同上同上layers.GlobalAveragePooling2D()

四、归一化层(Normalization Layers)

层名全称作用输入形状输出形状示例
BatchNormalization批归一化缩放标准化每一层输出同前一层输出同前一层输出layers.BatchNormalization()
LayerNormalization层归一化对每个样本单独归一化同前一层输出同前一层输出layers.LayerNormalization()
GroupNormalization组归一化分组后标准化同前一层输出同前一层输出自定义实现或使用 tfa.layers.GroupNormalization(来自 TF Addons)
InstanceNormalization实例归一化图像风格迁移常用同前一层输出同前一层输出tfa.layers.InstanceNormalization()

五、循环神经网络层(RNN Layers)

层名全称作用输入形状输出形状示例
SimpleRNN简单 RNN基础递归单元(batch_size, time_steps, input_dim)(batch_size, units)(batch_size, time_steps, units)layers.SimpleRNN(32)
LSTM长短期记忆解决长期依赖问题同上同上layers.LSTM(32)
GRU门控循环单元简化的 LSTM同上同上layers.GRU(32)
Bidirectional双向包装器封装 RNN 实现双向处理同 RNN输出维度翻倍layers.Bidirectional(layers.LSTM(32))
TimeDistributed时间分布封装器在时间维度上应用任意层(batch_size, time_steps, ...)(batch_size, time_steps, ...)layers.TimeDistributed(layers.Dense(10))

六、嵌入与查找层(Embedding Layers)

层名全称作用输入形状输出形状示例
Embedding嵌入层将索引映射为稠密向量(batch_size,)(batch_size, seq_len)(batch_size, output_dim)(batch_size, seq_len, output_dim)layers.Embedding(input_dim=10000, output_dim=64)
Hashing哈希层将字符串哈希到整数(batch_size, features)(batch_size, hash_buckets)tf.keras.layers.Hashing(num_bins=1000)
Discretization离散化层将连续值离散为区间索引(batch_size, features)(batch_size, features)tf.keras.layers.Discretization(bins=[10, 20, 30])

七、Transformer 相关层(Transformer Layers)

层名全称作用输入形状输出形状示例
MultiHeadAttention多头注意力机制实现自注意力机制(batch_size, seq_len, embed_dim)(batch_size, seq_len, embed_dim)layers.MultiHeadAttention(num_heads=8, key_dim=64)
PositionalEncoding位置编码添加位置信息(需自定义)同输入同输入自定义实现
TransformerTransformer 编码器/解码器包含多头注意力和 FFN(batch_size, seq_len, embed_dim)同输入自定义构建或使用 tfm.nlp.layers.TransformerEncoder(TF Model Garden)

八、展平与重塑层(Flatten & Reshape)

层名全称作用输入形状输出形状示例
Flatten展平操作将多维张量压成一维(batch_size, d1, d2, ..., dn)(batch_size, d1*d2*...*dn)layers.Flatten()
Reshape重塑操作改变张量形状(batch_size, d1, d2, ...)(batch_size, new_shape)layers.Reshape((28, 28, 1))
Permute维度重排改变维度顺序(batch_size, d1, d2, ...)(batch_size, reordered_dims)layers.Permute((2, 1))

九、高级层与工具层(Advanced / Utility Layers)

层名全称作用输入形状输出形状示例
Add张量相加实现残差连接多个张量(必须同形)同输入layers.Add()([x1, x2])
Multiply张量相乘实现逐元素相乘多个张量(必须同形)同输入layers.Multiply()([x1, x2])
Concatenate张量拼接沿指定轴拼接多个张量多个张量(除拼接轴外其他维度一致)新拼接张量layers.Concatenate(axis=-1)([x1, x2])
Lambda自定义函数封装包裹任意函数任意任意layers.Lambda(lambda x: x * 2)
Dot张量点积计算两个张量的点积(batch_size, dim1), (batch_size, dim2)(batch_size, 1)layers.Dot(axes=1)([x1, x2])

十、损失函数(Loss Functions)

不是”层”,但常用于模型编译

名称全称作用示例
MeanSquaredError均方误差损失回归任务losses.MeanSquaredError()
CategoricalCrossentropy多分类交叉熵分类任务losses.CategoricalCrossentropy()
BinaryCrossentropy二分类交叉熵二分类任务losses.BinaryCrossentropy()
SparseCategoricalCrossentropy稀疏多分类交叉熵整数标签分类任务losses.SparseCategoricalCrossentropy()
KLDivergenceKL 散度损失分布差异衡量losses.KLDivergence()

十一、正则化与约束(Regularizers & Constraints)

这些不是”层”,而是可以传入给层的参数。

名称类型说明示例
l1, l2, l1_l2正则化器权重衰减kernel_regularizer='l2'
NonNeg, MinMaxNorm, UnitNorm约束器权重限制kernel_constraint='non_neg'

十二、优化器(Optimizer)

优化器名称接口特点适用场景使用示例
SGD(随机梯度下降)tf.keras.optimizers.SGD(learning_rate=0.01, momentum=0.0, nesterov=False)最基础的优化方法;支持动量加速收敛广泛适用于各种问题,特别是大规模数据集optimizer = tf.keras.optimizers.SGD(learning_rate=0.01, momentum=0.9)
Adamtf.keras.optimizers.Adam(learning_rate=0.001, beta_1=0.9, beta_2=0.999, epsilon=1e-07)自适应学习率,结合了 AdaGrad 和 RMSProp 的优点大多数深度学习任务,默认选择之一optimizer = tf.keras.optimizers.Adam(learning_rate=0.001)
RMSproptf.keras.optimizers.RMSprop(learning_rate=0.001, rho=0.9, momentum=0.0, epsilon=1e-07)通过平方梯度的移动平均值来归一化学习率循环神经网络(RNN),非凸优化问题optimizer = tf.keras.optimizers.RMSprop(learning_rate=0.001, rho=0.9)
Adagradtf.keras.optimizers.Adagrad(learning_rate=0.001, initial_accumulator_value=0.1, epsilon=1e-07)学习率基于参数更新历史自适应调整稀疏数据,如自然语言处理(NLP)任务optimizer = tf.keras.optimizers.Adagrad(learning_rate=0.01)
Adadeltatf.keras.optimizers.Adadelta(learning_rate=0.001, rho=0.95, epsilon=1e-07)Adagrad 的改进版,减少对初始学习率的选择敏感性不需要手动设置学习率的情况optimizer = tf.keras.optimizers.Adadelta(learning_rate=1.0)
Adamaxtf.keras.optimizers.Adamax(learning_rate=0.001, beta_1=0.9, beta_2=0.999, epsilon=1e-07)Adam 的扩展版本,使用无穷范数更加稳定的学习过程optimizer = tf.keras.optimizers.Adamax(learning_rate=0.002)
Nadamtf.keras.optimizers.Nadam(learning_rate=0.001, beta_1=0.9, beta_2=0.999, epsilon=1e-07)结合 Nesterov 动量的 Adam高效且稳定的训练optimizer = tf.keras.optimizers.Nadam(learning_rate=0.002)
Ftrltf.keras.optimizers.Ftrl(learning_rate=0.1, learning_rate_power=-0.5, initial_accumulator_value=0.1, l1_regularization_strength=0.0, l2_regularization_strength=0.0)在线学习中表现出色,特别适合稀疏模型广告点击预测等稀疏数据场景optimizer = tf.keras.optimizers.Ftrl(learning_rate=0.1)

注意事项:

  • learning_rate:学习率决定了每一步权重更新的幅度。过大可能导致不收敛,过小则会使训练进程非常缓慢。
  • momentum/nesterov:动量技术帮助加速 SGD 在相关方向上的进展,并抑制震荡。Nesterov 动量是对标准动量的一种改进。
  • beta_1, beta_2, epsilon:这些参数主要用于 Adam 及其变体,控制了一阶矩估计的指数衰减速率、二阶矩估计的指数衰减速率以及数值稳定性的一个小常数。
  • weight_decay/l1_regularization_strength/l2_regularization_strength:正则化项用于防止过拟合。

十三、学习率调度器(Learning Rate Scheduler)

名称接口 / 实现方式特点适用场景使用示例
ConstantScheduletf.constant(value)学习率保持不变所有任务,默认不调整 LRlr = tf.constant(0.01)
PiecewiseConstantDecaytf.keras.optimizers.schedules.PiecewiseConstantDecay(boundaries, values)在指定的 step 或 epoch 处设置不同的学习率自定义多阶段学习率boundaries = [10000, 15000]; values = [1e-2, 1e-3, 1e-4]
ExponentialDecaytf.keras.optimizers.schedules.ExponentialDecay(initial_learning_rate, decay_steps, decay_rate)每隔一定步数按指数衰减 LR简单、平滑的 LR 下降策略lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(0.1, decay_steps=1000, decay_rate=0.9)
CosineDecaytf.keras.optimizers.schedules.CosineDecay(initial_learning_rate, decay_steps)按余弦退火方式降低 LR图像分类等任务lr_schedule = tf.keras.optimizers.schedules.CosineDecay(0.1, decay_steps=1000)
CosineDecayRestartstf.keras.optimizers.schedules.CosineDecayRestarts(initial_learning_rate, first_decay_steps, t_mul=2.0, m_mul=1.0, alpha=0.0)余弦退火 + 多周期重启需要周期性探索最优lr_schedule = tf.keras.optimizers.schedules.CosineDecayRestarts(0.1, 100)
InverseTimeDecaytf.keras.optimizers.schedules.InverseTimeDecay(initial_learning_rate, decay_steps, decay_rate)学习率随时间反比例下降慢速衰减,适合长训练lr_schedule = tf.keras.optimizers.schedules.InverseTimeDecay(0.01, 100, 0.5)
PolynomialDecaytf.keras.optimizers.schedules.PolynomialDecay(initial_learning_rate, decay_steps, end_learning_rate=0.0001, power=1.0)多项式形式从初始 LR 衰减到最小值可控的 LR 衰减过程lr_schedule = tf.keras.optimizers.schedules.PolynomialDecay(0.1, decay_steps=1000, end_learning_rate=0.001, power=2.0)
LearningRateSchedule(自定义)继承 tf.keras.optimizers.schedules.LearningRateSchedule 并重写 __call__() 方法完全自定义调度逻辑任意复杂需求见下方示例

如何选择:

目标推荐调度器
固定学习率tf.constant
分段设定 LRPiecewiseConstantDecay
简单指数衰减ExponentialDecay
余弦退火CosineDecay, CosineDecayRestarts
温和衰减InverseTimeDecay
控制最终 LRPolynomialDecay
自定义逻辑LearningRateSchedule 类继承实现

使用示例:

# 在 TensorFlow 中使用学习率调度器时,通常将它作为参数传入优化器:
optimizer = tf.keras.optimizers.Adam(learning_rate=lr_scheduler)
# 编译模型
model.compile(optimizer=optimizer, loss='sparse_categorical_crossentropy', metrics=['accuracy'])

注意

  • 如果你使用的是 tf.data.Dataset.fit() 训练流程,调度器会自动根据 step 更新学习率。
  • 如果是自定义训练循环(如 GradientTape),你需要手动管理 step 参数来更新学习率。

总结

类别数量
基础层(Core Layers)~4
卷积层(Convolutional Layers)~5
池化层(Pooling Layers)~4
归一化层(Normalization Layers)~4
RNN 层(RNN Layers)~4
嵌入层(Embedding Layers)~2
Transformer 层(Transformer Layers)~2
展平与重塑层(Flatten & Reshape)~3
高级操作层(Advanced / Utility Layers)~5
损失函数(Loss Functions)~5
正则化与约束(Regularizers & Constraints)~2

TensorFlow 中自定义层

1. 基础步骤

  • 继承 tf.keras.layers.Layer:所有自定义层都应该从这个基类继承。
  • 重写 __init__ 方法:用于初始化参数(如权重)和其他子层。
  • 重写 build 方法:定义权重。尽管可以在 __init__ 中定义权重,但使用 build 可以让你根据输入的形状动态地创建权重。
  • 重写 call 方法:定义前向传播逻辑。

2. 示例:创建一个简单的线性变换层

如果需要序列化或者保存模型,需要额外实现 get_config 方法。

下面的例子展示了如何创建一个简单的线性变换层,它执行的操作类似于 y = Wx + b

import tensorflow as tf

class SimpleDense(tf.keras.layers.Layer):
    def __init__(self, units=32, activation=None):
        super(SimpleDense, self).__init__()
        self.units = units
        self.activation = tf.keras.activations.get(activation)

    def build(self, input_shape):  # 创建权重
        # 添加权重矩阵
        self.w = self.add_weight(shape=(input_shape[-1], self.units),
                                 initializer='random_normal',
                                 trainable=True)
        # 添加偏置向量
        self.b = self.add_weight(shape=(self.units,),
                                 initializer='zeros',
                                 trainable=True)

    def call(self, inputs):  # 定义前向传播
        x = tf.matmul(inputs, self.w) + self.b
        if self.activation is not None:
            x = self.activation(x)
        return x

使用自定义层:

一旦你定义了自己的层,就可以像使用任何其他内置层一样使用它:

# 创建一个模型并使用自定义层
model = tf.keras.Sequential([
    SimpleDense(units=16, activation='relu'),
    SimpleDense(units=10)
])

# 编译模型
model.compile(optimizer='adam', loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True))

# 训练模型
data = tf.random.uniform((100, 3))
labels = tf.random.uniform((100,), maxval=10, dtype=tf.int32)
model.fit(data, labels, epochs=5)

3. 进阶示例:包含更多复杂性的自定义层

如果你想要创建更复杂的层,比如包含多个操作或者需要访问额外的配置选项,你可以扩展上述基本框架。例如,创建一个带有批归一化的简单卷积层:

class ConvBatchNorm(tf.keras.layers.Layer):
    def __init__(self, filters, kernel_size, strides=1, padding='valid'):
        super(ConvBatchNorm, self).__init__()
        self.filters = filters
        self.kernel_size = kernel_size
        self.strides = strides
        self.padding = padding

    def build(self, input_shape):
        self.conv = tf.keras.layers.Conv2D(filters=self.filters,
                                           kernel_size=self.kernel_size,
                                           strides=self.strides,
                                           padding=self.padding)
        self.bn = tf.keras.layers.BatchNormalization()

    def call(self, inputs, training=False):
        x = self.conv(inputs)
        return self.bn(x, training=training)

注意事项

  • 训练标志:某些层(如 BatchNormalizationDropout)的行为在训练和推理阶段不同。确保你的自定义层正确处理 training 参数。
  • 可学习参数:使用 self.add_weight() 方法添加可学习参数。
  • 动态形状支持:利用 build 方法可以根据实际输入数据的形状动态创建权重,从而支持不同大小的输入。

张量操作和连续存储

PyTorch 必须知悉,TensorFlow 自动处理

操作作用描述改变物理内存改变读取逻辑是否需要连续存储
tf.transpose()交换张量的维度顺序不需要
tf.reshape()改变张量的形状而不改变其数据可能(取决于原始张量的布局)TensorFlow 自动处理
tf.expand_dims()在指定维度插入一个新的维度不需要
tf.squeeze()移除大小为 1 的维度不需要
tf.slice()提取张量的一个切片不需要
tf.tile()复制张量沿每个维度不需要
tf.concat()沿指定轴连接张量列表不需要
tf.stack()按新轴堆叠一系列张量不需要
tf.split()沿指定轴分割张量不需要
tf.cast()转换张量的数据类型不需要
tf.identity()返回一个具有相同形状和内容的新张量不需要
tf.convert_to_tensor()将各种类型的输入转换为张量可能不需要

解释:

  • 改变物理内存:指该操作是否会实际导致数据被复制到新的内存位置。
  • 改变读取逻辑:表示该操作是否只是改变了数据的访问方式,而没有改变其在内存中的布局。
  • 是否需要前置内存整理操作:在 TensorFlow 中,虽然没有直接等同于 PyTorch 的 .contiguous() 函数,但在某些情况下,你可能需要确保张量是平坦或连续的以避免潜在的性能问题。不过,大多数时候 TensorFlow 会自动处理这些情况。

模型保存与加载

保存方法文件格式/目录结构是否包含模型结构是否推荐使用适用场景保存代码示例加载代码示例
model.save()model_savedmodel/ 目录,包含 saved_model.pb + variables/ 子目录✅(推荐)生产部署、跨平台调用(如 TensorFlow Serving)model.save('model_savedmodel')model = tf.keras.models.load_model('model_savedmodel')
model.save('model.h5').h5(HDF5 格式)✅(适合简单模型)快速保存模型用于训练恢复或小型项目model.save('model.h5')from tensorflow.keras.models import load_model; model = load_model('model.h5')
tf.saved_model.save()saved_model/ 目录(与上述 SavedModel 类似)✅(适用于更灵活的保存方式)高级保存需求、自定义签名等tf.saved_model.save(model, 'model_custom')model = tf.keras.models.load_model('model_custom')
model.to_json() + model.save_weights().json(结构)+ .h5(权重)✅ JSON 含结构,❌ 权重单独保存⚠️ 一般不推荐分开保存结构和权重见下方示例见下方示例
导出为 ONNX 格式.onnx✅(用于跨框架部署)跨框架迁移(如转 PyTorch、ONNX Runtime 等)pip install onnx-tf; onnx-tf convert -i model.h5 -o model.onnximport onnxruntime as ort; session = ort.InferenceSession("model.onnx")

补充说明

1. SavedModel(推荐)

  • 特点:
    • 包含完整模型信息(结构、权重、优化器状态等)
    • 支持版本控制
    • 可部署到 TensorFlow Serving、移动端、Web 等环境
  • 优点:
    • 官方推荐格式
    • 最适合生产环境使用

2. HDF5 (.h5)

  • 特点:
    • 单个文件,便于管理
    • 适用于简单的训练和推理流程
  • 缺点:
    • 不支持所有自定义层
    • 不如 SavedModel 强大

3. 分开保存结构和权重

  • 适用场景:
    • 需要对结构进行修改后再加载权重
    • 用于调试或特定研究需求
  • 注意:
    • 操作繁琐,容易出错,不建议作为默认方案
# 分开保存结构和权重示例
model_json = model.to_json()
with open("model.json", "w") as json_file:
    json_file.write(model_json)
model.save_weights("model_weights.h5")

# 加载
from tensorflow.keras.models import model_from_json
with open('model.json', 'r') as json_file:
    loaded_model_json = json_file.read()
loaded_model = model_from_json(loaded_model_json)
loaded_model.load_weights('model_weights.h5')

4. 导出为 ONNX

  • 用途:实现 TensorFlow 与其他框架(如 PyTorch、Caffe、ONNX Runtime)之间的互操作
  • 工具:onnx-tf 库支持将 TensorFlow 模型转换为 ONNX 格式