通用训练模板
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 调度器 |
| 自定义损失函数 | 替换 SparseCategoricalCrossentropy 为 FocalLoss、DiceLoss 等 |
| 模型评估与推理 | 加载 .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可视化训练过程
扩展建议
| 功能 | 描述 |
|---|---|
| 增加图像增强 | 添加 RandomFlip、RandomRotation 等数据增强操作 |
| 替换为原始 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可视化训练过程
扩展建议
| 功能 | 描述 |
|---|---|
| 增加图像增强 | 添加 RandomFlip、RandomRotation 等数据增强操作 |
| 替换为原始 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可视化训练过程
扩展建议
| 功能 | 描述 |
|---|---|
| 增加图像增强 | 添加 RandomFlip、RandomRotation 等数据增强操作 |
| 替换为原始 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可视化训练过程
扩展建议
| 功能 | 描述 |
|---|---|
| 增加图像增强 | 添加 RandomFlip、RandomRotation 等数据增强操作 |
| 替换为真实数据集 | 如医学影像数据集等 |
| 多 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)
| 部分 | 类型 | 参数说明 |
|---|---|---|
| Generator | Conv2DTranspose | 将噪声向量映射为真实图像 |
| Discriminator | Conv2D | 判断图像是真实还是生成的 |
| Loss Function | Hinge Loss | 更适合稳定训练 GAN |
| Optimizer | Adam(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可视化训练过程
扩展建议
| 功能 | 描述 |
|---|---|
| 增加图像增强 | 添加 RandomFlip、RandomRotation 等数据增强操作 |
| 替换为原始 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可视化训练过程
扩展建议
| 功能 | 描述 |
|---|---|
| 增加图像增强 | 添加 RandomFlip、RandomRotation 等数据增强操作 |
| 替换为原始 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可视化训练过程
扩展建议
| 功能 | 描述 |
|---|---|
| 增加图像增强 | 添加 RandomFlip、RandomRotation 等数据增强操作 |
| 替换为原始 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可视化训练过程
扩展建议
| 功能 | 描述 |
|---|---|
| 增加图像增强 | 添加 RandomFlip、RandomRotation 等数据增强操作 |
| 替换为原始 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可视化训练过程
扩展建议
| 功能 | 描述 |
|---|---|
| 增加图像增强 | 添加 RandomFlip、RandomRotation 等数据增强操作 |
| 替换为真实数据集 | 如医学影像数据集等 |
| 多 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)
| 部分 | 类型 | 参数说明 |
|---|---|---|
| Generator | Conv2DTranspose | 将噪声向量映射为真实图像 |
| Discriminator | Conv2D | 判断图像是真实还是生成的 |
| Loss Function | Hinge Loss | 更适合稳定训练 GAN |
| Optimizer | Adam(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可视化训练过程
扩展建议
| 功能 | 描述 |
|---|---|
| 增加图像增强 | 添加 RandomFlip、RandomRotation 等数据增强操作 |
| 替换为原始 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可视化训练过程
扩展建议
| 功能 | 描述 |
|---|---|
| 增加图像增强 | 添加 RandomFlip、RandomRotation 等数据增强操作 |
| 替换为原始 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可视化训练过程
扩展建议
| 功能 | 描述 |
|---|---|
| 增加图像增强 | 添加 RandomFlip、RandomRotation 等数据增强操作 |
| 替换为原始 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可视化训练过程
扩展建议
| 功能 | 描述 |
|---|---|
| 增加图像增强 | 添加 RandomFlip、RandomRotation 等数据增强操作 |
| 替换为真实数据集 | 如医学影像数据集等 |
| 多 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)
| 部分 | 类型 | 参数说明 |
|---|---|---|
| Generator | Conv2DTranspose | 将噪声向量映射为真实图像 |
| Discriminator | Conv2D | 判断图像是真实还是生成的 |
| Loss Function | Hinge Loss | 更适合稳定训练 GAN |
| Optimizer | Adam(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 Layer | MLP 层堆叠,融合所有特征 |
| 输出层 | 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 Function | CrossEntropyLoss 或 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 Function | CrossEntropyLoss 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)“,通过三个主要的”门”来控制信息的流动:
-
遗忘门(Forget Gate):决定哪些信息从 cell state(记忆单元)中被丢弃。
- 使用 sigmoid 激活函数输出一个 0~1 的值:
- 0 表示完全忘记
- 1 表示完全保留
- 使用 sigmoid 激活函数输出一个 0~1 的值:
-
输入门(Input Gate):决定哪些新信息将被加入到 cell state 中。
- 包括两个部分:
- 一个 sigmoid 层决定哪些值要更新
- 一个 tanh 层创建一个新的候选值向量
- 包括两个部分:
-
输出门(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/2D | layers.SeparableConv2D(32, (3,3), activation='relu') |
DepthwiseConv2D | 深度可分离卷积 | 单通道卷积 | 同 Conv2D | 同 Conv2D | layers.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 | 位置编码 | 添加位置信息(需自定义) | 同输入 | 同输入 | 自定义实现 |
Transformer | Transformer 编码器/解码器 | 包含多头注意力和 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() |
KLDivergence | KL 散度损失 | 分布差异衡量 | 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) |
| Adam | tf.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) |
| RMSprop | tf.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) |
| Adagrad | tf.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) |
| Adadelta | tf.keras.optimizers.Adadelta(learning_rate=0.001, rho=0.95, epsilon=1e-07) | Adagrad 的改进版,减少对初始学习率的选择敏感性 | 不需要手动设置学习率的情况 | optimizer = tf.keras.optimizers.Adadelta(learning_rate=1.0) |
| Adamax | tf.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) |
| Nadam | tf.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) |
| Ftrl | tf.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)
| 名称 | 接口 / 实现方式 | 特点 | 适用场景 | 使用示例 |
|---|---|---|---|---|
ConstantSchedule | tf.constant(value) | 学习率保持不变 | 所有任务,默认不调整 LR | lr = tf.constant(0.01) |
PiecewiseConstantDecay | tf.keras.optimizers.schedules.PiecewiseConstantDecay(boundaries, values) | 在指定的 step 或 epoch 处设置不同的学习率 | 自定义多阶段学习率 | boundaries = [10000, 15000]; values = [1e-2, 1e-3, 1e-4] |
ExponentialDecay | tf.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) |
CosineDecay | tf.keras.optimizers.schedules.CosineDecay(initial_learning_rate, decay_steps) | 按余弦退火方式降低 LR | 图像分类等任务 | lr_schedule = tf.keras.optimizers.schedules.CosineDecay(0.1, decay_steps=1000) |
CosineDecayRestarts | tf.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) |
InverseTimeDecay | tf.keras.optimizers.schedules.InverseTimeDecay(initial_learning_rate, decay_steps, decay_rate) | 学习率随时间反比例下降 | 慢速衰减,适合长训练 | lr_schedule = tf.keras.optimizers.schedules.InverseTimeDecay(0.01, 100, 0.5) |
PolynomialDecay | tf.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 |
| 分段设定 LR | PiecewiseConstantDecay |
| 简单指数衰减 | ExponentialDecay |
| 余弦退火 | CosineDecay, CosineDecayRestarts |
| 温和衰减 | InverseTimeDecay |
| 控制最终 LR | PolynomialDecay |
| 自定义逻辑 | 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)
注意事项
- 训练标志:某些层(如
BatchNormalization和Dropout)的行为在训练和推理阶段不同。确保你的自定义层正确处理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.onnx | import 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 格式