Article
通用训练模板
- 数据加载模块化（支持自定义数据集）
- 模型构建灵活（可替换任意网络结构）
- 支持 GPU 加速
- 完整训练 + 验证流程
- 日志记录和模型保存机制
- 可扩展性强（便于添加 LR 调度器、早停机制等）
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from torch.utils.tensorboard import SummaryWriter
import os
import logging
import time
# ================== Step 1: 自定义数据集 ==================
class MyDataset(Dataset):
def __init__(self):
# 示例：随机生成一些数据
self.data = torch.randn(1000, 3, 224, 224) # 示例图像数据 (batch_size, channels, H, W)
self.labels = torch.randint(0, 3, (1000,)) # 示例标签（三分类）
def __len__(self):
return len(self.data)
def __getitem__(self, idx):
return self.data[idx], self.labels[idx]
# ================== Step 2: 构建模型 ==================
class SimpleModel(nn.Module):
def __init__(self):
# 一定要调用super完成nn.Module的初始化工作，否则会导致以下问题：1.优化器找不到任何参数、2.模型无法保存和回复、3.无法使用to(device)，4.训练报错
super(SimpleModel, self).__init__()
self.feature_extractor = nn.Sequential(
nn.Conv2d(3, 16, kernel_size=3),
nn.ReLU(),
nn.MaxPool2d(2),
nn.Conv2d(16, 32, kernel_size=3),
nn.ReLU(),
nn.MaxPool2d(2),
)
self.classifier = nn.Sequential(
nn.Flatten(),
nn.Linear(32 * 54 * 54, 128),
nn.ReLU(),
nn.Linear(128, 3) # 最终输出维度为3维，未经过softmax处理的三分类结果
)
def forward(self, x):
x = self.feature_extractor(x)
x = self.classifier(x)
return x
# ================== Step 3: 工具函数与配置 ==================
def get_dataloaders(batch_size=32):
train_dataset = MyDataset()
val_dataset = MyDataset() # 实际中应使用真实验证集
# DataLoader选项：
# num_workers：使用多少个子进程加载数据，默认为0（主进程）
# prefetch_factor：每个worker预取多少个batch，默认为2
# pin_memory：是否将数据拷贝到CUDA pinned memory，加速GPU拷贝
# batch_size：每个batch的样本数量
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=batch_size)
return train_loader, val_loader
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 train_one_epoch(model, dataloader, criterion, optimizer, device, epoch, writer):
model.train()
running_loss = 0.0
correct = 0
total = 0
for i, (inputs, labels) in enumerate(dataloader):
inputs, labels = inputs.to(device), labels.to(device)
optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
running_loss += loss.item()
_, predicted = outputs.max(1)
total += labels.size(0)
correct += predicted.eq(labels).sum().item()
if i % 10 == 9:
avg_loss = running_loss / 10
acc = 100. * correct / total
logging.info(f"Epoch [{epoch}] Batch [{i}] Loss: {avg_loss:.4f}, Acc: {acc:.2f}%")
writer.add_scalar('Loss/train', avg_loss, epoch * len(dataloader) + i)
running_loss = 0.0
correct = 0
total = 0
@torch.no_grad()
def validate(model, dataloader, criterion, device, epoch, writer):
model.eval()
correct = 0
total = 0
val_loss = 0.0
for inputs, labels in dataloader:
inputs, labels = inputs.to(device), labels.to(device)
outputs = model(inputs)
loss = criterion(outputs, labels) # 使用long()或者float()对labels的tensor类型进行转型，例如labels.long()
val_loss += loss.item()
_, predicted = outputs.max(1)
total += labels.size(0)
correct += predicted.eq(labels).sum().item()
acc = 100. * correct / total
avg_loss = val_loss / len(dataloader)
logging.info(f"Epoch [{epoch}] Validation Acc: {acc:.2f}%, Loss: {avg_loss:.4f}")
writer.add_scalar('Loss/val', avg_loss, epoch)
writer.add_scalar('Accuracy/val', acc, epoch)
return acc
# ================== Step 5: 主程序入口 ==================
def main():
setup_logger()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
logging.info(f"Using device: {device}")
batch_size = 32
num_epochs = 10
lr = 0.001
train_loader, val_loader = get_dataloaders(batch_size)
model = SimpleModel().to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=lr)
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.1)
writer = SummaryWriter(log_dir="runs")
best_acc = 0.0
save_dir = "checkpoints"
os.makedirs(save_dir, exist_ok=True)
for epoch in range(num_epochs):
logging.info(f"\nEpoch [{epoch+1}/{num_epochs}]")
train_one_epoch(model, train_loader, criterion, optimizer, device, epoch+1, writer)
val_acc = validate(model, val_loader, criterion, device, epoch+1, writer)
scheduler.step()
# 保存最佳模型
if val_acc > best_acc:
best_acc = val_acc
torch.save(model.state_dict(), os.path.join(save_dir, f"best_model.pth"))
logging.info(f"Saved best model with accuracy: {best_acc:.2f}%")
writer.close()
logging.info("Training complete.")
if __name__ == "__main__":
main()LeNet
为了使用提供的通用训练模板来实现 LeNet 模型，我们需要对模板进行一些调整以适应 LeNet 的结构和 MNIST 数据集。LeNet 主要用于手写数字识别任务（MNIST 数据集），因此我们将使用 torchvision.datasets.MNIST 来加载数据，并定义 LeNet 模型架构。
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
from torchvision import datasets, transforms
import os
import logging
import time
# ================== Step 1: 加载并预处理 MNIST 数据集 ==================
def get_dataloaders(batch_size=64):
transform = transforms.Compose([
transforms.ToTensor(), # 将PIL图像或Numpy数组转换为PyTorch张量Tensor，将像素值从[0, 255]归一化到[0,1]
transforms.Normalize((0.1307,), (0.3081,)) # 对图像进行标准化操作，对每个通道减去均值0.1307，除以标准差0.3081
])
train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
val_dataset = datasets.MNIST(root='./data', train=False, transform=transform)
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=batch_size)
return train_loader, val_loader
# ================== Step 2: 定义 LeNet 模型 ==================
class LeNet(nn.Module):
def __init__(self):
super(LeNet, self).__init__()
self.feature_extractor = nn.Sequential(
nn.Conv2d(1, 6, kernel_size=5), # 输入通道数为1，输出通道数为6
nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2),
nn.Conv2d(6, 16, kernel_size=5),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2)
)
self.classifier = nn.Sequential(
nn.Linear(16 * 4 * 4, 120),
nn.ReLU(),
nn.Linear(120, 84),
nn.ReLU(),
nn.Linear(84, 10) # MNIST有10个类别
)
def forward(self, x):
x = self.feature_extractor(x)
x = torch.flatten(x, 1) # 展平操作
x = self.classifier(x)
return x
# ================== 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 train_one_epoch(model, dataloader, criterion, optimizer, device, epoch, writer):
model.train()
running_loss = 0.0
correct = 0
total = 0
for i, (inputs, labels) in enumerate(dataloader):
inputs, labels = inputs.to(device), labels.to(device)
optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
running_loss += loss.item()
_, predicted = outputs.max(1)
total += labels.size(0)
correct += predicted.eq(labels).sum().item()
if i % 10 == 9:
avg_loss = running_loss / 10
acc = 100. * correct / total
logging.info(f"Epoch [{epoch}] Batch [{i}] Loss: {avg_loss:.4f}, Acc: {acc:.2f}%")
writer.add_scalar('Loss/train', avg_loss, epoch * len(dataloader) + i)
running_loss = 0.0
correct = 0
total = 0
@torch.no_grad()
def validate(model, dataloader, criterion, device, epoch, writer):
model.eval()
correct = 0
total = 0
val_loss = 0.0
for inputs, labels in dataloader:
inputs, labels = inputs.to(device), labels.to(device)
outputs = model(inputs)
loss = criterion(outputs, labels)
val_loss += loss.item()
_, predicted = outputs.max(1)
total += labels.size(0)
correct += predicted.eq(labels).sum().item()
acc = 100. * correct / total
avg_loss = val_loss / len(dataloader)
logging.info(f"Epoch [{epoch}] Validation Acc: {acc:.2f}%, Loss: {avg_loss:.4f}")
writer.add_scalar('Loss/val', avg_loss, epoch)
writer.add_scalar('Accuracy/val', acc, epoch)
return acc
# ================== Step 5: 主程序入口 ==================
def main():
setup_logger()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
logging.info(f"Using device: {device}")
batch_size = 64
num_epochs = 10
lr = 0.001
train_loader, val_loader = get_dataloaders(batch_size)
model = LeNet().to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=lr)
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.1)
writer = SummaryWriter(log_dir="runs")
best_acc = 0.0
save_dir = "checkpoints"
os.makedirs(save_dir, exist_ok=True)
for epoch in range(num_epochs):
logging.info(f"\nEpoch [{epoch+1}/{num_epochs}]")
train_one_epoch(model, train_loader, criterion, optimizer, device, epoch+1, writer)
val_acc = validate(model, val_loader, criterion, device, epoch+1, writer)
scheduler.step()
# 保存最佳模型
if val_acc > best_acc:
best_acc = val_acc
torch.save(model.state_dict(), os.path.join(save_dir, f"best_model.pth"))
logging.info(f"Saved best model with accuracy: {best_acc:.2f}%")
writer.close()
logging.info("Training complete.")
if __name__ == "__main__":
main()关键点说明
- 数据加载：使用
torchvision.datasets.MNIST加载 MNIST 数据集，应用了标准的图像归一化 (mean=0.1307,std=0.3081)。 - LeNet 模型定义：包含两个卷积层，每个卷积层后跟一个 ReLU 激活函数和最大池化层。全连接层部分包括三个线性层，最终输出为 10 维向量（对应于 10 个类别）。
- 训练流程：标准的训练循环，包括前向传播、损失计算、反向传播及参数更新。每隔一定步数打印日志信息，并记录到 TensorBoard 中。验证阶段评估模型性能，并保存最优模型。
通过以下命令启动 TensorBoard 查看训练过程中的指标：
# runs是日志所在文件夹
tensorboard --logdir runsAlexNet
为了使用提供的通用训练模板来实现 AlexNet 模型，我们需要对模板进行一些调整以适应 AlexNet 的结构和 ImageNet 数据集。由于 ImageNet 数据集非常大，通常不适合直接下载到本地用于示例，因此我们将使用 CIFAR-10 数据集作为替代，它也是一个图像分类数据集，并且大小适中，适合快速实验。
import torch
from torch.utils.data import DataLoader
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.tensorboard import SummaryWriter
import os
import logging
import time
# ================== Step 1: 加载并预处理 CIFAR-10 数据集 ==================
def get_dataloaders(batch_size=64):
transform_train = transforms.Compose([
transforms.Resize((227, 227)), # AlexNet输入尺寸
transforms.RandomHorizontalFlip(), # 以给定的概率随机水平翻转图像
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
transform_val = transforms.Compose([
transforms.Resize((227, 227)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
train_dataset = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform_train)
val_dataset = datasets.CIFAR10(root='./data', train=False, transform=transform_val)
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=batch_size)
return train_loader, val_loader
# ================== Step 2: 定义 AlexNet 模型 ==================
class AlexNet(nn.Module):
def __init__(self, num_classes=10):
super(AlexNet, self).__init__()
self.features = nn.Sequential(
nn.Conv2d(3, 64, kernel_size=11, stride=4, padding=2),
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=3, stride=2),
nn.Conv2d(64, 192, kernel_size=5, padding=2),
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=3, stride=2),
nn.Conv2d(192, 384, kernel_size=3, padding=1),
nn.ReLU(inplace=True),
nn.Conv2d(384, 256, kernel_size=3, padding=1),
nn.ReLU(inplace=True),
nn.Conv2d(256, 256, kernel_size=3, padding=1),
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=3, stride=2),
)
self.classifier = nn.Sequential(
nn.Dropout(),
nn.Linear(256 * 6 * 6, 4096),
nn.ReLU(inplace=True),
nn.Dropout(),
nn.Linear(4096, 4096),
nn.ReLU(inplace=True),
nn.Linear(4096, num_classes),
)
def forward(self, x):
x = self.features(x)
x = torch.flatten(x, 1)
x = self.classifier(x)
return x
# ================== 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 train_one_epoch(model, dataloader, criterion, optimizer, device, epoch, writer):
model.train()
running_loss = 0.0
correct = 0
total = 0
for i, (inputs, labels) in enumerate(dataloader):
inputs, labels = inputs.to(device), labels.to(device)
optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
running_loss += loss.item()
_, predicted = outputs.max(1)
total += labels.size(0)
correct += predicted.eq(labels).sum().item()
if i % 10 == 9:
avg_loss = running_loss / 10
acc = 100. * correct / total
logging.info(f"Epoch [{epoch}] Batch [{i}] Loss: {avg_loss:.4f}, Acc: {acc:.2f}%")
writer.add_scalar('Loss/train', avg_loss, epoch * len(dataloader) + i)
running_loss = 0.0
correct = 0
total = 0
@torch.no_grad()
def validate(model, dataloader, criterion, device, epoch, writer):
model.eval()
correct = 0
total = 0
val_loss = 0.0
for inputs, labels in dataloader:
inputs, labels = inputs.to(device), labels.to(device)
outputs = model(inputs)
loss = criterion(outputs, labels)
val_loss += loss.item()
_, predicted = outputs.max(1)
total += labels.size(0)
correct += predicted.eq(labels).sum().item()
acc = 100. * correct / total
avg_loss = val_loss / len(dataloader)
logging.info(f"Epoch [{epoch}] Validation Acc: {acc:.2f}%, Loss: {avg_loss:.4f}")
writer.add_scalar('Loss/val', avg_loss, epoch)
writer.add_scalar('Accuracy/val', acc, epoch)
return acc
# ================== Step 5: 主程序入口 ==================
def main():
setup_logger()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
logging.info(f"Using device: {device}")
batch_size = 64
num_epochs = 10
lr = 0.001
train_loader, val_loader = get_dataloaders(batch_size)
model = AlexNet(num_classes=10).to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=lr)
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.1)
writer = SummaryWriter(log_dir="runs")
best_acc = 0.0
save_dir = "checkpoints"
os.makedirs(save_dir, exist_ok=True)
for epoch in range(num_epochs):
logging.info(f"\nEpoch [{epoch+1}/{num_epochs}]")
train_one_epoch(model, train_loader, criterion, optimizer, device, epoch+1, writer)
val_acc = validate(model, val_loader, criterion, device, epoch+1, writer)
scheduler.step()
# 保存最佳模型
if val_acc > best_acc:
best_acc = val_acc
torch.save(model.state_dict(), os.path.join(save_dir, f"best_model.pth"))
logging.info(f"Saved best model with accuracy: {best_acc:.2f}%")
writer.close()
logging.info("Training complete.")
if __name__ == "__main__":
main()关键点说明
- 数据加载：使用
torchvision.datasets.CIFAR10加载 CIFAR-10 数据集。对图像进行了 Resize 操作使其符合 AlexNet 输入要求（227x227），同时应用了标准化处理。 - AlexNet 模型定义：包含五个卷积层，每个卷积层后跟一个 ReLU 激活函数以及最大池化层（除了最后的卷积层）。全连接层部分包括三个线性层，前两个全连接层后面跟着 Dropout 层以防止过拟合。
- 训练流程：标准的训练循环，包括前向传播、损失计算、反向传播及参数更新。每隔一定步数打印日志信息，并记录到 TensorBoard 中。验证阶段评估模型性能，并保存最优模型。
安装必要的库：
pip install torchvision tensorboard通过以下命令启动 TensorBoard 查看训练过程中的指标：
tensorboard --logdir runsVGG
为了使用提供的通用训练模板实现 VGG 模型，我们将以 VGG16 为例，并在 CIFAR-10 数据集上进行训练。VGG 模型以其简单而有效的架构设计著称，通过堆叠多个 3x3 的小卷积核来增加网络深度。
import torch
from torch.utils.data import DataLoader
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.tensorboard import SummaryWriter
import os
import logging
import time
# ================== Step 1: 加载并预处理 CIFAR-10 数据集 ==================
def get_dataloaders(batch_size=64):
transform_train = transforms.Compose([
transforms.Resize((224, 224)), # VGG输入尺寸
transforms.RandomHorizontalFlip(), # 以给定的概率随机水平翻转图像
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
transform_val = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
train_dataset = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform_train)
val_dataset = datasets.CIFAR10(root='./data', train=False, transform=transform_val)
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=batch_size)
return train_loader, val_loader
# ================== Step 2: 定义 VGG16 模型 ==================
class VGG16(nn.Module):
def __init__(self, num_classes=10):
super(VGG16, self).__init__()
self.features = nn.Sequential(
nn.Conv2d(3, 64, kernel_size=3, padding=1),
nn.ReLU(inplace=True),
nn.Conv2d(64, 64, kernel_size=3, padding=1),
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=2, stride=2),
nn.Conv2d(64, 128, kernel_size=3, padding=1),
nn.ReLU(inplace=True),
nn.Conv2d(128, 128, kernel_size=3, padding=1),
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=2, stride=2),
nn.Conv2d(128, 256, kernel_size=3, padding=1),
nn.ReLU(inplace=True),
nn.Conv2d(256, 256, kernel_size=3, padding=1),
nn.ReLU(inplace=True),
nn.Conv2d(256, 256, kernel_size=3, padding=1),
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=2, stride=2),
nn.Conv2d(256, 512, kernel_size=3, padding=1),
nn.ReLU(inplace=True),
nn.Conv2d(512, 512, kernel_size=3, padding=1),
nn.ReLU(inplace=True),
nn.Conv2d(512, 512, kernel_size=3, padding=1),
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=2, stride=2),
nn.Conv2d(512, 512, kernel_size=3, padding=1),
nn.ReLU(inplace=True),
nn.Conv2d(512, 512, kernel_size=3, padding=1),
nn.ReLU(inplace=True),
nn.Conv2d(512, 512, kernel_size=3, padding=1),
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=2, stride=2),
)
self.classifier = nn.Sequential(
nn.Linear(512 * 7 * 7, 4096),
nn.ReLU(True),
nn.Dropout(),
nn.Linear(4096, 4096),
nn.ReLU(True),
nn.Dropout(),
nn.Linear(4096, num_classes),
)
def forward(self, x):
x = self.features(x)
x = torch.flatten(x, 1)
x = self.classifier(x)
return x
# ================== 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 train_one_epoch(model, dataloader, criterion, optimizer, device, epoch, writer):
model.train()
running_loss = 0.0
correct = 0
total = 0
for i, (inputs, labels) in enumerate(dataloader):
inputs, labels = inputs.to(device), labels.to(device)
optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
running_loss += loss.item()
_, predicted = outputs.max(1)
total += labels.size(0)
correct += predicted.eq(labels).sum().item()
if i % 10 == 9:
avg_loss = running_loss / 10
acc = 100. * correct / total
logging.info(f"Epoch [{epoch}] Batch [{i}] Loss: {avg_loss:.4f}, Acc: {acc:.2f}%")
writer.add_scalar('Loss/train', avg_loss, epoch * len(dataloader) + i)
running_loss = 0.0
correct = 0
total = 0
@torch.no_grad()
def validate(model, dataloader, criterion, device, epoch, writer):
model.eval()
correct = 0
total = 0
val_loss = 0.0
for inputs, labels in dataloader:
inputs, labels = inputs.to(device), labels.to(device)
outputs = model(inputs)
loss = criterion(outputs, labels)
val_loss += loss.item()
_, predicted = outputs.max(1)
total += labels.size(0)
correct += predicted.eq(labels).sum().item()
acc = 100. * correct / total
avg_loss = val_loss / len(dataloader)
logging.info(f"Epoch [{epoch}] Validation Acc: {acc:.2f}%, Loss: {avg_loss:.4f}")
writer.add_scalar('Loss/val', avg_loss, epoch)
writer.add_scalar('Accuracy/val', acc, epoch)
return acc
# ================== Step 5: 主程序入口 ==================
def main():
setup_logger()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
logging.info(f"Using device: {device}")
batch_size = 64
num_epochs = 10
lr = 0.001
train_loader, val_loader = get_dataloaders(batch_size)
model = VGG16(num_classes=10).to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=lr)
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.1)
writer = SummaryWriter(log_dir="runs")
best_acc = 0.0
save_dir = "checkpoints"
os.makedirs(save_dir, exist_ok=True)
for epoch in range(num_epochs):
logging.info(f"\nEpoch [{epoch+1}/{num_epochs}]")
train_one_epoch(model, train_loader, criterion, optimizer, device, epoch+1, writer)
val_acc = validate(model, val_loader, criterion, device, epoch+1, writer)
scheduler.step()
# 保存最佳模型
if val_acc > best_acc:
best_acc = val_acc
torch.save(model.state_dict(), os.path.join(save_dir, f"best_model.pth"))
logging.info(f"Saved best model with accuracy: {best_acc:.2f}%")
writer.close()
logging.info("Training complete.")
if __name__ == "__main__":
main()关键点说明
- 数据加载：使用
torchvision.datasets.CIFAR10加载 CIFAR-10 数据集。对图像进行了 Resize 操作使其符合 VGG 输入要求（224x224），同时应用了标准化处理。 - VGG16 模型定义：包含 13 个卷积层和 3 个全连接层，每个卷积层后跟一个 ReLU 激活函数以及最大池化层（除了最后的卷积层）。全连接层部分包括三个线性层，前两个全连接层后面跟着 Dropout 层以防止过拟合。
- 训练流程：标准的训练循环，包括前向传播、损失计算、反向传播及参数更新。每隔一定步数打印日志信息，并记录到 TensorBoard 中。验证阶段评估模型性能，并保存最优模型。
安装必要的库：
pip install torchvision tensorboard通过以下命令启动 TensorBoard 查看训练过程中的指标：
tensorboard --logdir runsGoogLeNet / Inception
要使用通用训练模板实现 GoogLeNet / Inception v1，我们需要：
- 使用 Inception 模块构建网络结构
- 在 CIFAR-10 数据集上进行训练（因为 ImageNet 太大）
- 保持与之前一致的训练流程：数据加载、模型定义、训练/验证循环、日志记录、模型保存等
GoogLeNet（也叫 Inception v1）是 Google 提出的一种高效 CNN 架构，核心思想是：
- 使用 Inception 模块
- 并行使用不同大小的卷积核（1x1, 3x3, 5x5），提升特征提取能力
- 使用 1x1 卷积降维减少参数量和计算量
import torch
from torch.utils.data import DataLoader
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.tensorboard import SummaryWriter
import os
import logging
import time
# ================== Step 1: 加载并预处理 CIFAR-10 数据集 ==================
def get_dataloaders(batch_size=64):
transform_train = transforms.Compose([
transforms.Resize((224, 224)), # GoogLeNet 输入尺寸为 224x224
transforms.RandomHorizontalFlip(), # 以给定的概率随机水平翻转图像
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
transform_val = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
train_dataset = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform_train)
val_dataset = datasets.CIFAR10(root='./data', train=False, transform=transform_val)
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=batch_size)
return train_loader, val_loader
# ================== Step 2: 定义 Inception 模块 ==================
class Inception(nn.Module):
def __init__(self, in_channels, ch1x1, ch3x3red, ch3x3, ch5x5red, ch5x5, pool_proj):
super(Inception, self).__init__()
self.branch1 = nn.Conv2d(in_channels, ch1x1, kernel_size=1)
self.branch2 = nn.Sequential(
nn.Conv2d(in_channels, ch3x3red, kernel_size=1),
nn.Conv2d(ch3x3red, ch3x3, kernel_size=3, padding=1)
)
self.branch3 = nn.Sequential(
nn.Conv2d(in_channels, ch5x5red, kernel_size=1),
nn.Conv2d(ch5x5red, ch5x5, kernel_size=5, padding=2)
)
self.branch4 = nn.Sequential(
nn.MaxPool2d(kernel_size=3, stride=1, padding=1),
nn.Conv2d(in_channels, pool_proj, kernel_size=1)
)
def forward(self, x):
branch1 = self.branch1(x)
branch2 = self.branch2(x)
branch3 = self.branch3(x)
branch4 = self.branch4(x)
return torch.cat([branch1, branch2, branch3, branch4], dim=1)
# ================== Step 3: 定义完整 GoogLeNet 模型 ==================
class GoogLeNet(nn.Module):
def __init__(self, num_classes=10):
super(GoogLeNet, self).__init__()
self.conv1 = nn.Sequential(
nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3),
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
)
self.conv2 = nn.Sequential(
nn.Conv2d(64, 64, kernel_size=1),
nn.ReLU(inplace=True),
nn.Conv2d(64, 192, kernel_size=3, padding=1),
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
)
self.inception3a = Inception(192, 64, 96, 128, 16, 32, 32)
self.inception3b = Inception(256, 128, 128, 192, 32, 96, 64)
self.max_pool3 = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
self.inception4a = Inception(480, 192, 96, 208, 16, 48, 64)
self.inception4b = Inception(512, 160, 112, 224, 24, 64, 64)
self.inception4c = Inception(512, 128, 128, 256, 24, 64, 64)
self.inception4d = Inception(512, 112, 144, 288, 32, 64, 64)
self.inception4e = Inception(528, 256, 160, 320, 32, 128, 128)
self.max_pool4 = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
self.inception5a = Inception(832, 256, 160, 320, 32, 128, 128)
self.inception5b = Inception(832, 384, 192, 384, 48, 128, 128)
self.avg_pool = nn.AdaptiveAvgPool2d((1, 1))
self.dropout = nn.Dropout(0.4)
self.fc = nn.Linear(1024, num_classes)
def forward(self, x):
x = self.conv1(x)
x = self.conv2(x)
x = self.inception3a(x)
x = self.inception3b(x)
x = self.max_pool3(x)
x = self.inception4a(x)
x = self.inception4b(x)
x = self.inception4c(x)
x = self.inception4d(x)
x = self.inception4e(x)
x = self.max_pool4(x)
x = self.inception5a(x)
x = self.inception5b(x)
x = self.avg_pool(x)
x = self.dropout(x)
x = torch.flatten(x, 1)
x = self.fc(x)
return x
# ================== 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 train_one_epoch(model, dataloader, criterion, optimizer, device, epoch, writer):
model.train()
running_loss = 0.0
correct = 0
total = 0
for i, (inputs, labels) in enumerate(dataloader):
inputs, labels = inputs.to(device), labels.to(device)
optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
running_loss += loss.item()
_, predicted = outputs.max(1)
total += labels.size(0)
correct += predicted.eq(labels).sum().item()
if i % 10 == 9:
avg_loss = running_loss / 10
acc = 100. * correct / total
logging.info(f"Epoch [{epoch}] Batch [{i}] Loss: {avg_loss:.4f}, Acc: {acc:.2f}%")
writer.add_scalar('Loss/train', avg_loss, epoch * len(dataloader) + i)
running_loss = 0.0
correct = 0
total = 0
@torch.no_grad()
def validate(model, dataloader, criterion, device, epoch, writer):
model.eval()
correct = 0
total = 0
val_loss = 0.0
for inputs, labels in dataloader:
inputs, labels = inputs.to(device), labels.to(device)
outputs = model(inputs)
loss = criterion(outputs, labels)
val_loss += loss.item()
_, predicted = outputs.max(1)
total += labels.size(0)
correct += predicted.eq(labels).sum().item()
acc = 100. * correct / total
avg_loss = val_loss / len(dataloader)
logging.info(f"Epoch [{epoch}] Validation Acc: {acc:.2f}%, Loss: {avg_loss:.4f}")
writer.add_scalar('Loss/val', avg_loss, epoch)
writer.add_scalar('Accuracy/val', acc, epoch)
return acc
# ================== Step 6: 主程序入口 ==================
def main():
setup_logger()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
logging.info(f"Using device: {device}")
batch_size = 32
num_epochs = 10
lr = 0.001
train_loader, val_loader = get_dataloaders(batch_size)
model = GoogLeNet(num_classes=10).to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=lr)
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.1)
writer = SummaryWriter(log_dir="runs")
best_acc = 0.0
save_dir = "checkpoints"
os.makedirs(save_dir, exist_ok=True)
for epoch in range(num_epochs):
logging.info(f"\nEpoch [{epoch+1}/{num_epochs}]")
train_one_epoch(model, train_loader, criterion, optimizer, device, epoch+1, writer)
val_acc = validate(model, val_loader, criterion, device, epoch+1, writer)
scheduler.step()
# 保存最佳模型
if val_acc > best_acc:
best_acc = val_acc
torch.save(model.state_dict(), os.path.join(save_dir, f"best_model.pth"))
logging.info(f"Saved best model with accuracy: {best_acc:.2f}%")
writer.close()
logging.info("Training complete.")
if __name__ == "__main__":
main()关键说明
| 部分 | 内容 |
|---|---|
| 模型结构 | 使用 Inception 模块构建完整的 GoogLeNet（Inception v1） |
| 输入尺寸 | Resize 到 224x224，符合 GoogLeNet 的输入要求 |
| 分类类别数 | num_classes=10（适用于 CIFAR-10） |
| 优化器 | Adam + StepLR 学习率调度器 |
| 日志记录 | 使用 TensorBoard 和 Logging 输出训练过程 |
| 模型保存 | 保存最佳准确率对应的模型权重 |
安装必要的依赖库：
pip install torch torchvision tensorboard启动 TensorBoard 查看训练日志：
tensorboard --logdir runs扩展建议
- 添加辅助分类器（GoogLeNet 原始设计中有两个辅助输出）
- 替换为更高级的 Inception 变体（如 Inception v3/v4）
- 使用 ImageNet 数据集进行大规模训练（需 GPU 资源）
- 使用混合精度训练加速收敛（
torch.cuda.amp） - 添加早停机制（Early Stopping）
ResNet
ResNet（残差网络）是深度学习中非常经典且高效的卷积神经网络结构，它通过引入残差连接（skip connection / shortcut connection）解决了深层网络中的梯度消失问题，使得模型可以训练得更深。
我们将使用 PyTorch 实现一个完整的 ResNet-18 模型，并基于 CIFAR-10 数据集进行训练。代码结构将遵循”通用训练模板”，包括：
- 数据加载与预处理
- 模型定义（ResNet-18）
- 工具函数（日志、模型保存等）
- 训练和验证函数
- 主程序入口
项目目标
| 目标项 | 内容 |
|---|---|
| 模型 | ResNet-18 |
| 数据集 | CIFAR-10 |
| 输入尺寸 | 32x32 → Resize 到 224x224（适配 ResNet）或保持 32x32（轻量版） |
| 损失函数 | CrossEntropyLoss |
| 优化器 | Adam + StepLR 学习率调度 |
| 日志记录 | TensorBoard + Logging |
| 模型保存 | 最佳模型保存 |
ResNet 的核心思想
ResNet 的关键在于残差块（Residual Block）：
class ResidualBlock(nn.Module):
def __init__(self, in_channels):
super(ResidualBlock, self).__init__()
self.conv1 = nn.Conv2d(in_channels, in_channels, kernel_size=3, padding=1)
self.bn1 = nn.BatchNorm2d(in_channels)
self.conv2 = nn.Conv2d(in_channels, in_channels, kernel_size=3, padding=1)
self.bn2 = nn.BatchNorm2d(in_channels)
def forward(self, x):
residual = x
x = torch.relu(self.bn1(self.conv1(x)))
x = self.bn2(self.conv2(x))
x += residual
return torch.relu(x)完整代码示例（ResNet-18）
import torch
from torch.utils.data import DataLoader
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.tensorboard import SummaryWriter
import os
import logging
import time
# ================== Step 1: 加载并预处理 CIFAR-10 数据集 ==================
def get_dataloaders(batch_size=64):
transform_train = transforms.Compose([
transforms.Resize((224, 224)), # ResNet 输入为 224x224
transforms.RandomHorizontalFlip(), # 以给定的概率随机水平翻转图像
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
transform_val = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
train_dataset = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform_train)
val_dataset = datasets.CIFAR10(root='./data', train=False, transform=transform_val)
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=batch_size)
return train_loader, val_loader
# ================== Step 2: 定义 ResNet-18 模型 ==================
class BasicBlock(nn.Module):
expansion = 1
def __init__(self, inplanes, planes, stride=1, downsample=None):
super(BasicBlock, self).__init__()
self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=3, stride=stride, padding=1, bias=False)
self.bn1 = nn.BatchNorm2d(planes)
self.relu = nn.ReLU(inplace=True)
self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, padding=1, bias=False)
self.bn2 = nn.BatchNorm2d(planes)
self.downsample = downsample
self.stride = stride
def forward(self, x):
identity = x
out = self.conv1(x)
out = self.bn1(out)
out = self.relu(out)
out = self.conv2(out)
out = self.bn2(out)
if self.downsample is not None:
identity = self.downsample(x)
out += identity
out = self.relu(out)
return out
class ResNet(nn.Module):
def __init__(self, block, layers, num_classes=10):
super(ResNet, self).__init__()
self.inplanes = 64
self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False)
self.bn1 = nn.BatchNorm2d(64)
self.relu = nn.ReLU(inplace=True)
self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
self.layer1 = self._make_layer(block, 64, layers[0])
self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
self.fc = nn.Linear(512 * block.expansion, num_classes)
def _make_layer(self, block, planes, blocks, stride=1):
downsample = None
if stride != 1 or self.inplanes != planes * block.expansion:
downsample = nn.Sequential(
nn.Conv2d(self.inplanes, planes * block.expansion, kernel_size=1, stride=stride, bias=False),
nn.BatchNorm2d(planes * block.expansion),
)
layers = []
layers.append(block(self.inplanes, planes, stride, downsample))
self.inplanes = planes * block.expansion
for _ in range(1, blocks):
layers.append(block(self.inplanes, planes))
return nn.Sequential(*layers)
def forward(self, x):
x = self.conv1(x)
x = self.bn1(x)
x = self.relu(x)
x = self.maxpool(x)
x = self.layer1(x)
x = self.layer2(x)
x = self.layer3(x)
x = self.layer4(x)
x = self.avgpool(x)
x = torch.flatten(x, 1)
x = self.fc(x)
return x
def ResNet18():
return ResNet(BasicBlock, [2, 2, 2, 2], num_classes=10)
# ================== 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 train_one_epoch(model, dataloader, criterion, optimizer, device, epoch, writer):
model.train()
running_loss = 0.0
correct = 0
total = 0
for i, (inputs, labels) in enumerate(dataloader):
inputs, labels = inputs.to(device), labels.to(device)
optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
running_loss += loss.item()
_, predicted = outputs.max(1)
total += labels.size(0)
correct += predicted.eq(labels).sum().item()
if i % 10 == 9:
avg_loss = running_loss / 10
acc = 100. * correct / total
logging.info(f"Epoch [{epoch}] Batch [{i}] Loss: {avg_loss:.4f}, Acc: {acc:.2f}%")
writer.add_scalar('Loss/train', avg_loss, epoch * len(dataloader) + i)
running_loss = 0.0
correct = 0
total = 0
@torch.no_grad()
def validate(model, dataloader, criterion, device, epoch, writer):
model.eval()
correct = 0
total = 0
val_loss = 0.0
for inputs, labels in dataloader:
inputs, labels = inputs.to(device), labels.to(device)
outputs = model(inputs)
loss = criterion(outputs, labels)
val_loss += loss.item()
_, predicted = outputs.max(1)
total += labels.size(0)
correct += predicted.eq(labels).sum().item()
acc = 100. * correct / total
avg_loss = val_loss / len(dataloader)
logging.info(f"Epoch [{epoch}] Validation Acc: {acc:.2f}%, Loss: {avg_loss:.4f}")
writer.add_scalar('Loss/val', avg_loss, epoch)
writer.add_scalar('Accuracy/val', acc, epoch)
return acc
# ================== Step 5: 主程序入口 ==================
def main():
setup_logger()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
logging.info(f"Using device: {device}")
batch_size = 32
num_epochs = 10
lr = 0.001
train_loader, val_loader = get_dataloaders(batch_size)
model = ResNet18().to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=lr)
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.1)
writer = SummaryWriter(log_dir="runs")
best_acc = 0.0
save_dir = "checkpoints"
os.makedirs(save_dir, exist_ok=True)
for epoch in range(num_epochs):
logging.info(f"\nEpoch [{epoch+1}/{num_epochs}]")
train_one_epoch(model, train_loader, criterion, optimizer, device, epoch+1, writer)
val_acc = validate(model, val_loader, criterion, device, epoch+1, writer)
scheduler.step()
# 保存最佳模型
if val_acc > best_acc:
best_acc = val_acc
torch.save(model.state_dict(), os.path.join(save_dir, f"best_model.pth"))
logging.info(f"Saved best model with accuracy: {best_acc:.2f}%")
writer.close()
logging.info("Training complete.")
if __name__ == "__main__":
main()关键说明
| 模块 | 功能 |
|---|---|
BasicBlock | ResNet-18 使用的基本残差块 |
ResNet18() | 构建完整 ResNet-18 模型 |
transforms.Resize((224, 224)) | 使输入符合 ImageNet 尺寸要求 |
nn.AdaptiveAvgPool2d((1, 1)) | 确保输出维度固定，适配不同输入尺寸 |
SummaryWriter | 可视化训练过程（损失、准确率） |
logging | 输出训练信息到控制台和日志文件 |
安装依赖：
pip install torch torchvision tensorboard启动 TensorBoard 查看训练日志：
tensorboard --logdir runs扩展建议
- 替换为 ResNet-34 / 50 / 101（只需修改 block 和 layer 数量）
- 使用混合精度训练加速收敛（
torch.cuda.amp） - 添加早停机制（Early Stopping）
- 在 ImageNet 上训练（需调整数据加载器）
U-Net
U-Net 是一种经典的编码器-解码器结构的卷积神经网络，广泛应用于图像分割任务中（如医学图像分割、遥感图像处理等）。它由一个下采样路径（收缩路径）和一个上采样路径（扩展路径）组成，并通过跳跃连接（skip connection）将不同层级的特征图进行拼接。
项目目标
| 目标项 | 内容 |
|---|---|
| 模型 | U-Net |
| 数据集 | 示例使用自制数据集或 torchvision.datasets 中的数据 |
| 输入尺寸 | 256x256 或其他适配尺寸 |
| 损失函数 | BCEWithLogitsLoss（适用于二分类分割） |
| 优化器 | Adam + StepLR 学习率调度 |
| 日志记录 | TensorBoard + Logging |
| 模型保存 | 最佳模型保存 |
U-Net 核心结构
U-Net 包含两个主要部分：
- 下采样路径（Encoder）：多个卷积块 + MaxPooling
- 上采样路径（Decoder）：上采样操作 + 卷积块 + 跳跃连接拼接
每个卷积块通常包括两个 Conv2d + BatchNorm + ReLU。
完整代码示例
我们将以二值图像分割任务为例，使用自制数据集格式（image, mask），并遵循”通用训练模板”。
import torch
from torch.utils.data import Dataset, DataLoader
import torch.nn as nn
import torch.optim as optim
from torch.utils.tensorboard import SummaryWriter
import os
import logging
import time
from torchvision import transforms
from PIL import Image
import glob
import numpy as np
# ================== Step 1: 自定义图像分割数据集 ==================
class SegmentationDataset(Dataset):
def __init__(self, image_dir, mask_dir, transform=None):
self.image_paths = sorted(glob.glob(os.path.join(image_dir, "*.png")))
self.mask_paths = sorted(glob.glob(os.path.join(mask_dir, "*.png")))
self.transform = transform
def __len__(self):
return len(self.image_paths)
def __getitem__(self, idx):
image = Image.open(self.image_paths[idx]).convert("L") # 灰度图
mask = Image.open(self.mask_paths[idx]).convert("L") # 灰度掩码
image = np.array(image).astype(np.float32) / 255.0
mask = np.array(mask).astype(np.float32) / 255.0
image = torch.from_numpy(image).unsqueeze(0)
mask = torch.from_numpy(mask).unsqueeze(0)
if self.transform:
image = self.transform(image)
return image, mask
def get_dataloaders(batch_size=4):
transform = transforms.Compose([
transforms.Resize((256, 256)),
])
train_dataset = SegmentationDataset(
image_dir='./data/train/images',
mask_dir='./data/train/masks',
transform=transform
)
val_dataset = SegmentationDataset(
image_dir='./data/val/images',
mask_dir='./data/val/masks',
transform=transform
)
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=batch_size)
return train_loader, val_loader
# ================== Step 2: 定义 U-Net 模型 ==================
class UNet(nn.Module):
def __init__(self, in_channels=1, out_channels=1):
super(UNet, self).__init__()
def double_conv(in_channels, out_channels):
return nn.Sequential(
nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
nn.BatchNorm2d(out_channels),
nn.ReLU(inplace=True),
nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
nn.BatchNorm2d(out_channels),
nn.ReLU(inplace=True)
)
self.down1 = double_conv(in_channels, 64)
self.down2 = double_conv(64, 128)
self.down3 = double_conv(128, 256)
self.down4 = double_conv(256, 512)
self.maxpool = nn.MaxPool2d(kernel_size=2)
self.up3 = nn.ConvTranspose2d(512, 256, kernel_size=2, stride=2)
self.up_conv3 = double_conv(512, 256)
self.up2 = nn.ConvTranspose2d(256, 128, kernel_size=2, stride=2)
self.up_conv2 = double_conv(256, 128)
self.up1 = nn.ConvTranspose2d(128, 64, kernel_size=2, stride=2)
self.up_conv1 = double_conv(128, 64)
self.final_conv = nn.Conv2d(64, out_channels, kernel_size=1)
def forward(self, x):
# 编码器
conv1 = self.down1(x)
x = self.maxpool(conv1)
conv2 = self.down2(x)
x = self.maxpool(conv2)
conv3 = self.down3(x)
x = self.maxpool(conv3)
x = self.down4(x)
# 解码器
x = self.up3(x)
x = torch.cat([x, conv3], dim=1)
x = self.up_conv3(x)
x = self.up2(x)
x = torch.cat([x, conv2], dim=1)
x = self.up_conv2(x)
x = self.up1(x)
x = torch.cat([x, conv1], dim=1)
x = self.up_conv1(x)
return self.final_conv(x)
# ================== 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 train_one_epoch(model, dataloader, criterion, optimizer, device, epoch, writer):
model.train()
running_loss = 0.0
for i, (inputs, targets) in enumerate(dataloader):
inputs, targets = inputs.to(device), targets.to(device)
optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, targets)
loss.backward()
optimizer.step()
running_loss += loss.item()
if i % 10 == 9:
avg_loss = running_loss / 10
logging.info(f"Epoch [{epoch}] Batch [{i}] Loss: {avg_loss:.4f}")
writer.add_scalar('Loss/train', avg_loss, epoch * len(dataloader) + i)
running_loss = 0.0
@torch.no_grad()
def validate(model, dataloader, criterion, device, epoch, writer):
model.eval()
total_loss = 0.0
total_iou = 0.0
for inputs, targets in dataloader:
inputs, targets = inputs.to(device), targets.to(device)
outputs = model(inputs)
loss = criterion(outputs, targets)
total_loss += loss.item()
preds = torch.sigmoid(outputs) > 0.5
intersection = (preds & targets.byte()).float().sum()
union = (preds | targets.byte()).float().sum()
iou = intersection / (union + 1e-8)
total_iou += iou.item()
avg_loss = total_loss / len(dataloader)
avg_iou = total_iou / len(dataloader)
logging.info(f"Epoch [{epoch}] Validation Loss: {avg_loss:.4f}, IoU: {avg_iou:.4f}")
writer.add_scalar('Loss/val', avg_loss, epoch)
writer.add_scalar('IoU/val', avg_iou, epoch)
return avg_iou
# ================== Step 5: 主程序入口 ==================
def main():
setup_logger()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
logging.info(f"Using device: {device}")
batch_size = 4
num_epochs = 20
lr = 1e-3
train_loader, val_loader = get_dataloaders(batch_size)
model = UNet().to(device)
criterion = nn.BCEWithLogitsLoss()
optimizer = optim.Adam(model.parameters(), lr=lr)
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.5)
writer = SummaryWriter(log_dir="runs")
best_iou = 0.0
save_dir = "checkpoints"
os.makedirs(save_dir, exist_ok=True)
for epoch in range(num_epochs):
logging.info(f"\nEpoch [{epoch+1}/{num_epochs}]")
train_one_epoch(model, train_loader, criterion, optimizer, device, epoch+1, writer)
val_iou = validate(model, val_loader, criterion, device, epoch+1, writer)
scheduler.step()
# 保存最佳模型
if val_iou > best_iou:
best_iou = val_iou
torch.save(model.state_dict(), os.path.join(save_dir, f"best_model.pth"))
logging.info(f"Saved best model with IoU: {best_iou:.4f}")
writer.close()
logging.info("Training complete.")
if __name__ == "__main__":
main()关键说明
| 部分 | 内容 |
|---|---|
| 数据集 | 自定义 SegmentationDataset 类，读取图像和对应 mask |
| 模型结构 | U-Net，包含跳跃连接和上下采样模块 |
| 损失函数 | BCEWithLogitsLoss（用于二分类分割） |
| 评估指标 | IoU（Intersection over Union） |
| 模型保存 | 仅保留 IoU 最高的模型 |
| 扩展性 | 可支持多类分割（修改输出通道数 + 使用 CrossEntropyLoss） |
安装依赖：
pip install torch torchvision pillow tensorboard启动 TensorBoard 查看训练日志：
tensorboard --logdir runs扩展建议
- 支持多类别分割（修改输出通道数 + 使用 CrossEntropyLoss）
- 添加可视化功能（TensorBoard 显示输入图像和预测 mask）
- 使用预训练编码器（如 ResNet 作为 encoder）
- 加入混合精度训练（
torch.cuda.amp） - 增加早停机制（Early Stopping）
GANs
生成对抗网络（GANs, Generative Adversarial Networks）是一种深度学习模型，由两个部分组成：生成器（Generator）和判别器（Discriminator）。生成器试图生成看起来真实的样本以欺骗判别器，而判别器则尝试区分真实样本与生成样本。通过这种对抗过程，生成器能够学习到数据的分布并生成新的、类似训练数据的样本。
项目目标
| 模块 | 内容 |
|---|---|
| 模型类型 | Vanilla GAN |
| 数据集 | MNIST（可替换为任意图像数据集） |
| 损失函数 | BCEWithLogitsLoss |
| 优化器 | Adam |
| 日志工具 | TensorBoard + Logging |
| 模型保存 | 最佳模型保存机制 |
| 扩展性 | 支持更换网络结构、损失函数等 |
目录结构建议
gan_template/
├── models/ # 模型定义
│ └── gan.py # Generator & Discriminator
├── datasets/ # 数据加载与预处理
│ └── mnist_dataset.py
├── utils/ # 工具类和辅助函数
│ ├── logger.py
│ └── saver.py
├── configs/ # 配置文件
│ └── config.yaml
├── main.py # 主程序入口
└── runs/ # TensorBoard 日志目录
└── images/ # 生成图像保存路径完整代码示例
import torch
from torch import nn, optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
from torch.utils.tensorboard import SummaryWriter
import os
import logging
import time
from torchvision.utils import make_grid
# ================== Step 1: 配置参数 ==================
class Config:
def __init__(self):
self.batch_size = 64
self.latent_dim = 100
self.lr = 2e-4
self.epochs = 50
self.device = "cuda" if torch.cuda.is_available() else "cpu"
self.log_dir = "runs/gan"
self.checkpoint_dir = "checkpoints"
self.image_save_dir = "runs/images"
self.n_save_samples = 64
config = Config()
# ================== Step 2: 定义 Generator 和 Discriminator ==================
class Generator(nn.Module):
def __init__(self, latent_dim=100):
super(Generator, self).__init__()
self.model = nn.Sequential(
nn.Linear(latent_dim, 128),
nn.LeakyReLU(0.2, inplace=True),
nn.Linear(128, 256),
nn.BatchNorm1d(256),
nn.LeakyReLU(0.2, inplace=True),
nn.Linear(256, 512),
nn.BatchNorm1d(512),
nn.LeakyReLU(0.2, inplace=True),
nn.Linear(512, 784),
nn.Tanh()
)
def forward(self, z):
img = self.model(z)
return img.view(img.size(0), 1, 28, 28)
class Discriminator(nn.Module):
def __init__(self):
super(Discriminator, self).__init__()
self.model = nn.Sequential(
nn.Linear(784, 512),
nn.LeakyReLU(0.2, inplace=True),
nn.Linear(512, 256),
nn.LeakyReLU(0.2, inplace=True),
nn.Linear(256, 1),
)
def forward(self, img):
validity = self.model(img.view(img.size(0), -1))
return validity
# ================== Step 3: 加载数据集 ==================
def get_dataloader(batch_size):
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5,), (0.5,))
])
dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
loader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
return loader
# ================== 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 save_checkpoint(state, path="checkpoints"):
os.makedirs(path, exist_ok=True)
torch.save(state, os.path.join(path, f"best_model.pth"))
# ================== Step 5: 训练函数 ==================
def train_one_epoch(loader, generator, discriminator, optimizer_g, optimizer_d, criterion, device, epoch, writer):
generator.train()
discriminator.train()
for i, (real_imgs, _) in enumerate(loader):
real_imgs = real_imgs.to(device)
batch_size = real_imgs.shape[0]
# 真实标签和假标签
valid = torch.ones((batch_size, 1), device=device)
fake = torch.zeros((batch_size, 1), device=device)
# ---------------------
# 训练 Generator
# ---------------------
optimizer_g.zero_grad()
z = torch.randn(batch_size, config.latent_dim, device=device)
gen_imgs = generator(z)
g_loss = criterion(discriminator(gen_imgs), valid)
g_loss.backward()
optimizer_g.step()
# ---------------------
# 训练 Discriminator
# ---------------------
optimizer_d.zero_grad()
real_loss = criterion(discriminator(real_imgs), valid)
fake_loss = criterion(discriminator(gen_imgs.detach()), fake)
d_loss = (real_loss + fake_loss) / 2
d_loss.backward()
optimizer_d.step()
# ---------------------
# 日志记录
# ---------------------
if i % 10 == 0:
logging.info(f"Epoch [{epoch}] Batch [{i}] D Loss: {d_loss.item():.4f}, G Loss: {g_loss.item():.4f}")
writer.add_scalar('Loss/Discriminator', d_loss.item(), epoch * len(loader) + i)
writer.add_scalar('Loss/Generator', g_loss.item(), epoch * len(loader) + i)
with torch.no_grad():
sample = generator(torch.randn(config.n_save_samples, config.latent_dim, device=device))
grid = make_grid(sample, nrow=8, normalize=True)
writer.add_image('Generated Images', grid, epoch * len(loader) + i)
return d_loss.item(), g_loss.item()
# ================== Step 6: 主程序入口 ==================
def main():
setup_logger(config.log_dir)
os.makedirs(config.image_save_dir, exist_ok=True)
logging.info("Loading dataset...")
dataloader = get_dataloader(config.batch_size)
logging.info("Initializing model...")
generator = Generator(config.latent_dim).to(config.device)
discriminator = Discriminator().to(config.device)
criterion = nn.BCEWithLogitsLoss()
optimizer_g = optim.Adam(generator.parameters(), lr=config.lr, betas=(0.5, 0.999))
optimizer_d = optim.Adam(discriminator.parameters(), lr=config.lr, betas=(0.5, 0.999))
writer = SummaryWriter(log_dir=config.log_dir)
best_g_loss = float('inf')
logging.info("Start training...")
for epoch in range(1, config.epochs + 1):
logging.info(f"\nEpoch [{epoch}/{config.epochs}]")
d_loss, g_loss = train_one_epoch(dataloader, generator, discriminator, optimizer_g, optimizer_d,
criterion, config.device, epoch, writer)
if g_loss < best_g_loss:
best_g_loss = g_loss
save_checkpoint({
'generator': generator.state_dict(),
'discriminator': discriminator.state_dict(),
'optimizer_g': optimizer_g.state_dict(),
'optimizer_d': optimizer_d.state_dict(),
'epoch': epoch,
'g_loss': g_loss,
'd_loss': d_loss,
}, config.checkpoint_dir)
logging.info(f"Saved best model with G Loss: {g_loss:.4f}")
writer.close()
logging.info("Training complete.")
if __name__ == "__main__":
main()关键说明
| 模块 | 功能 |
|---|---|
| Config 类 | 集中式配置管理，便于后期迁移至 YAML 文件 |
| Generator 和 Discriminator | 可替换为你自己的网络结构（如 DCGAN） |
| get_dataloader() | 可适配其他图像数据集 |
| train_one_epoch() | 模块化训练逻辑，支持 TensorBoard 可视化 |
| save_checkpoint() | 支持恢复训练或推理使用 |
| logging + SummaryWriter | 同时输出到终端和 TensorBoard |
确保已安装依赖：
pip install torch torchvision tensorboard Pillow启动 TensorBoard 查看训练过程：
tensorboard --logdir runs扩展建议
| 功能 | 描述 |
|---|---|
| 更复杂的网络 | 替换为 DCGAN、StyleGAN、WGAN-GP 等 |
| 多种损失函数 | 添加 Wasserstein 损失、GP 梯度惩罚等 |
| 图像可视化 | 增加生成图像的保存与展示功能 |
| 多 GPU 支持 | 使用 torch.nn.DataParallel |
| 断点续训 | 从 checkpoint 中恢复训练状态 |
| 测试模式 | 添加 generate_samples() 函数用于推理阶段 |
Transformer
Transformer 是一种基于自注意力机制（Self-Attention）的深度学习模型，最初提出用于自然语言处理任务（如机器翻译），后广泛应用于图像、语音等多个领域。
下面我们将实现一个完整的 Transformer 模型，并以序列到序列（Seq2Seq）任务为例（如英文到法语的翻译）。该模板具有良好的模块化设计，支持扩展为图像分类、文本摘要、问答系统等任务。
项目目标
| 模块 | 内容 |
|---|---|
| 模型类型 | Transformer（Encoder + Decoder） |
| 任务类型 | 序列到序列（Seq2Seq） |
| 数据集 | 自定义文本对数据集或 torchtext.datasets 中的翻译数据集 |
| 损失函数 | CrossEntropyLoss |
| 优化器 | Adam + StepLR 学习率调度 |
| 日志记录 | Logging + TensorBoard |
| 模型保存 | Checkpoint 机制 |
| 可扩展性 | 支持图像、文本摘要、多任务等 |
Transformer 核心结构
Transformer 包含两个主要部分：
- Encoder（编码器）
- 多头自注意力（Multi-head Self-Attention）
- 前馈网络（FeedForward）
- 层归一化（LayerNorm）
- 位置编码（Positional Encoding）
- Decoder（解码器）
- 掩码多头自注意力（Masked Multi-head Self-Attention）
- 编码器-解码器注意力（Cross Attention）
- 前馈网络
- 位置编码
- Transformer Model
- 组合 Encoder 和 Decoder
- 最终输出通过线性层映射为词表大小
完整代码示例
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
import torch.optim as optim
from torch.nn.utils.rnn import pad_sequence
from torch.utils.tensorboard import SummaryWriter
import os
import logging
import time
# ================== Step 1: 数据预处理与自定义数据集 ==================
class TranslationDataset(Dataset):
def __init__(self, data, src_tokenizer, tgt_tokenizer, max_len=50):
self.data = data
self.src_tokenizer = src_tokenizer
self.tgt_tokenizer = tgt_tokenizer
self.max_len = max_len
def __len__(self):
return len(self.data)
def __getitem__(self, idx):
src_sentence, tgt_sentence = self.data[idx]
# Tokenize and convert to indices
src_ids = self.src_tokenizer.encode(src_sentence)
tgt_ids = self.tgt_tokenizer.encode(tgt_sentence)
return torch.tensor(src_ids), torch.tensor(tgt_ids)
def collate_fn(batch):
src_batch, tgt_batch = zip(*batch)
src_batch = pad_sequence(src_batch, padding_value=0, batch_first=True)
tgt_batch = pad_sequence(tgt_batch, padding_value=0, batch_first=True)
return src_batch, tgt_batch
class SimpleTokenizer:
def __init__(self):
self.vocab = {"<pad>": 0, "<sos>": 1, "<eos>": 2}
self.reversed_vocab = {v: k for k, v in self.vocab.items()}
self.next_id = 3
def add_sentence(self, sentence):
tokens = sentence.split()
for token in tokens:
if token not in self.vocab:
self.vocab[token] = self.next_id
self.reversed_vocab[self.next_id] = token
self.next_id += 1
def encode(self, sentence):
tokens = sentence.split()
ids = [self.vocab.get(token, 0) for token in tokens]
return [1] + ids + [2] # Add SOS and EOS
def decode(self, ids):
return " ".join([self.reversed_vocab.get(i.item(), "<unk>") for i in ids if i not in (0, 1, 2)])
def get_dataloaders(batch_size=64):
# 示例数据：你可以替换为你自己的平行语料库
pairs = [
("I love NLP", "J'adore le traitement du langage naturel"),
("Hello world", "Bonjour le monde"),
("How are you", "Comment allez-vous"),
("Machine learning is fun", "Le machine learning c'est amusant"),
]
src_tokenizer = SimpleTokenizer()
tgt_tokenizer = SimpleTokenizer()
for eng, fra in pairs:
src_tokenizer.add_sentence(eng)
tgt_tokenizer.add_sentence(fra)
dataset = TranslationDataset(pairs, src_tokenizer, tgt_tokenizer)
loader = DataLoader(dataset, batch_size=batch_size, collate_fn=collate_fn, shuffle=True)
return loader, src_tokenizer, tgt_tokenizer
# ================== Step 2: 实现 Transformer 模型 ==================
class PositionalEncoding(nn.Module):
def __init__(self, d_model, max_len=5000):
super(PositionalEncoding, self).__init__()
pe = torch.zeros(max_len, d_model)
position = torch.arange(0, max_len).unsqueeze(1)
div_term = torch.exp(torch.arange(0, d_model, 2) * -(torch.log(torch.tensor(10000.0)) / d_model))
pe[:, 0::2] = torch.sin(position * div_term)
pe[:, 1::2] = torch.cos(position * div_term)
pe = pe.unsqueeze(0)
self.register_buffer('pe', pe)
def forward(self, x):
x = x + self.pe[:, :x.size(1)]
return x
class TransformerModel(nn.Module):
def __init__(self, src_vocab_size, tgt_vocab_size, d_model=512, num_heads=8, num_layers=3, dim_feedforward=2048, dropout=0.1):
super().__init__()
self.embedding_dim = d_model
self.encoder_embedding = nn.Embedding(src_vocab_size, d_model)
self.decoder_embedding = nn.Embedding(tgt_vocab_size, d_model)
self.positional_encoding = PositionalEncoding(d_model)
self.transformer = nn.Transformer(
d_model=d_model,
nhead=num_heads,
num_encoder_layers=num_layers,
num_decoder_layers=num_layers,
dim_feedforward=dim_feedforward,
dropout=dropout
)
self.fc_out = nn.Linear(d_model, tgt_vocab_size)
self.d_model = d_model
def make_mask(self, size):
mask = torch.triu(torch.ones(size, size) == 1).transpose(0, 1)
mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
return mask
def forward(self, src, tgt):
src_key_padding_mask = (src == 0)
tgt_key_padding_mask = (tgt == 0)
src = self.encoder_embedding(src) * torch.sqrt(torch.tensor(self.d_model, dtype=torch.float32))
src = self.positional_encoding(src)
tgt = self.decoder_embedding(tgt) * torch.sqrt(torch.tensor(self.d_model, dtype=torch.float32))
tgt = self.positional_encoding(tgt)
tgt_mask = self.make_mask(tgt.size(1)).to(src.device)
transformer_output = self.transformer(
src.transpose(0, 1),
tgt.transpose(0, 1),
src_key_padding_mask=src_key_padding_mask,
tgt_mask=tgt_mask,
memory_key_padding_mask=src_key_padding_mask
)
output = self.fc_out(transformer_output.transpose(0, 1))
return output
# ================== 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 train_one_epoch(model, dataloader, criterion, optimizer, device, epoch, writer):
model.train()
running_loss = 0.0
for i, (src, tgt) in enumerate(dataloader):
src, tgt = src.to(device), tgt.to(device)
# Shift targets for teacher forcing
input_tgt = tgt[:, :-1]
target_tgt = tgt[:, 1:]
optimizer.zero_grad()
output = model(src, input_tgt)
loss = criterion(output.reshape(-1, output.shape[-1]), target_tgt.reshape(-1))
loss.backward()
optimizer.step()
running_loss += loss.item()
if i % 10 == 9:
avg_loss = running_loss / 10
logging.info(f"Epoch [{epoch}] Batch [{i}] Loss: {avg_loss:.4f}")
writer.add_scalar('Loss/train', avg_loss, epoch * len(dataloader) + i)
running_loss = 0.0
@torch.no_grad()
def validate(model, dataloader, criterion, device, epoch, writer):
model.eval()
total_loss = 0.0
for src, tgt in dataloader:
src, tgt = src.to(device), tgt.to(device)
input_tgt = tgt[:, :-1]
target_tgt = tgt[:, 1:]
output = model(src, input_tgt)
loss = criterion(output.reshape(-1, output.shape[-1]), target_tgt.reshape(-1))
total_loss += loss.item()
avg_loss = total_loss / len(dataloader)
logging.info(f"Epoch [{epoch}] Validation Loss: {avg_loss:.4f}")
writer.add_scalar('Loss/val', avg_loss, epoch)
return avg_loss
# ================== Step 5: 主程序入口 ==================
def main():
setup_logger()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
logging.info(f"Using device: {device}")
batch_size = 4
num_epochs = 50
d_model = 512
num_heads = 8
num_layers = 3
lr = 1e-3
train_loader, src_tokenizer, tgt_tokenizer = get_dataloaders(batch_size)
model = TransformerModel(
src_vocab_size=len(src_tokenizer.vocab),
tgt_vocab_size=len(tgt_tokenizer.vocab),
d_model=d_model,
num_heads=num_heads,
num_layers=num_layers
).to(device)
criterion = nn.CrossEntropyLoss(ignore_index=0)
optimizer = optim.Adam(model.parameters(), lr=lr)
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.5)
writer = SummaryWriter(log_dir="runs")
best_val_loss = float('inf')
save_dir = "checkpoints"
os.makedirs(save_dir, exist_ok=True)
for epoch in range(1, num_epochs + 1):
logging.info(f"\nEpoch [{epoch}/{num_epochs}]")
train_one_epoch(model, train_loader, criterion, optimizer, device, epoch, writer)
val_loss = validate(model, train_loader, criterion, device, epoch, writer)
scheduler.step()
if val_loss < best_val_loss:
best_val_loss = val_loss
torch.save({
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
'epoch': epoch,
'loss': val_loss,
}, os.path.join(save_dir, f"best_model.pth"))
logging.info(f"Saved best model with Val Loss: {best_val_loss:.4f}")
writer.close()
logging.info("Training complete.")
if __name__ == "__main__":
main()关键说明
| 模块 | 功能 |
|---|---|
| SimpleTokenizer | 简单的分词器，可替换为 BPE、WordPiece、SpaCy、transformers 等 |
| TranslationDataset | 支持任意格式的平行语料库 |
| collate_fn | 对变长序列进行填充 |
| PositionalEncoding | 添加位置信息，使 Transformer 能理解顺序 |
| TransformerModel | 使用 PyTorch 的 nn.Transformer 模块构建完整模型 |
| train_one_epoch | 模块化的训练逻辑 |
| validate | 验证逻辑 |
| logging + SummaryWriter | 同时输出日志和 TensorBoard 可视化 |
| save_checkpoint | 保存最佳模型状态 |
确保安装以下依赖：
pip install torch tensorboard启动 TensorBoard 查看训练过程：
tensorboard --logdir runs扩展建议
| 功能 | 描述 |
|---|---|
| 更复杂的数据集 | 替换为 WMT、IWSLT 或你自己的语料库 |
| 更强的 Tokenizer | 使用 transformers 提供的 BPE、BERT tokenizer |
| 混合精度训练 | 使用 torch.cuda.amp 提升训练效率 |
| 多 GPU 支持 | 使用 DataParallel 或 DistributedDataParallel |
| 图像任务 | 将 Transformer 改为 Vision Transformer（ViT） |
| 推理模式 | 添加 generate() 函数用于推理阶段 |
BERT
BERT（Bidirectional Encoder Representations from Transformers）是一种基于 Transformer 编码器的预训练语言模型，能够捕捉文本中的双向上下文信息。它在 NLP 领域具有广泛的应用，如文本分类、命名实体识别、问答系统等。
下面我们将实现一个完整的 BERT 模型，并以文本分类任务为例（如 IMDb 情感分析）。该模板结构清晰、模块化良好，支持扩展为多种下游任务。
项目目标
| 模块 | 内容 |
|---|---|
| 模型类型 | BERT（基于 Hugging Face Transformers） |
| 任务类型 | 文本分类 |
| 数据集 | IMDb / AG News / 自定义文本数据集 |
| 损失函数 | CrossEntropyLoss |
| 优化器 | AdamW + LinearWithWarmup 学习率调度 |
| 日志记录 | Logging + TensorBoard |
| 模型保存 | Checkpoint 机制 |
| 可扩展性 | 支持序列标注、问答、多任务等 |
BERT 核心结构（HuggingFace 实现）
我们使用 HuggingFace Transformers 提供的 BertModel 和 BertTokenizer，其核心结构如下：
- Embedding 层
- Token Embeddings
- Segment Embeddings
- Positional Embeddings
- Transformer Encoder 堆叠层
- 多头自注意力（Multi-head Self-Attention）
- 前馈网络（FeedForward）
- LayerNorm + Dropout
- Pooler 层
- 用于提取 [CLS] 向量，适配分类任务
完整代码示例
import torch
from torch import nn, optim
from torch.utils.data import Dataset, DataLoader
from torch.utils.tensorboard import SummaryWriter
from transformers import BertTokenizer, BertModel
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score
import os
import logging
import time
import pandas as pd
# ================== Step 1: 配置参数 ==================
class Config:
def __init__(self):
self.model_name = 'bert-base-uncased'
self.max_length = 128
self.batch_size = 16
self.lr = 2e-5
self.epochs = 5
self.device = "cuda" if torch.cuda.is_available() else "cpu"
self.log_dir = "runs/bert"
self.checkpoint_dir = "checkpoints"
self.num_classes = 2 # 修改为你自己的类别数
self.warmup_steps = 500
self.weight_decay = 0.01
config = Config()
# ================== Step 2: 加载并处理数据集 ==================
def load_dataset():
# 示例：加载 IMDb 数据集（可替换为任意 CSV 文件）
df = pd.read_csv('imdb.csv') # 列名应为 text 和 label
texts = df['text'].tolist()
labels = df['label'].tolist()
return texts, labels
class TextDataset(Dataset):
def __init__(self, texts, labels, tokenizer, max_length=128):
self.texts = texts
self.labels = labels
self.tokenizer = tokenizer
self.max_length = max_length
def __len__(self):
return len(self.texts)
def __getitem__(self, idx):
text = self.texts[idx]
label = self.labels[idx]
encoding = self.tokenizer(
text,
add_special_tokens=True,
max_length=self.max_length,
padding='max_length',
truncation=True,
return_tensors='pt'
)
return {
'input_ids': encoding['input_ids'].flatten(),
'attention_mask': encoding['attention_mask'].flatten(),
'label': torch.tensor(label, dtype=torch.long)
}
def get_dataloaders(batch_size=16):
texts, labels = load_dataset()
train_texts, val_texts, train_labels, val_labels = train_test_split(texts, labels, test_size=0.2, random_state=42)
tokenizer = BertTokenizer.from_pretrained(config.model_name)
train_dataset = TextDataset(train_texts, train_labels, tokenizer, config.max_length)
val_dataset = TextDataset(val_texts, val_labels, tokenizer, config.max_length)
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False)
return train_loader, val_loader, tokenizer
# ================== Step 3: 定义 BERT 分类模型 ==================
class BERTClassifier(nn.Module):
def __init__(self, model_name, num_classes):
super(BERTClassifier, self).__init__()
self.bert = BertModel.from_pretrained(model_name)
self.classifier = nn.Linear(self.bert.config.hidden_size, num_classes)
def forward(self, input_ids, attention_mask):
outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)
cls_output = outputs.pooler_output # [CLS] 向量
logits = self.classifier(cls_output)
return logits
# ================== 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 save_checkpoint(state, path="checkpoints"):
os.makedirs(path, exist_ok=True)
torch.save(state, os.path.join(path, f"best_model.pth"))
# ================== Step 5: 训练与验证函数 ==================
def train_one_epoch(model, dataloader, criterion, optimizer, scheduler, device, epoch, writer):
model.train()
running_loss = 0.0
all_preds, all_labels = [], []
for i, batch in enumerate(dataloader):
input_ids = batch['input_ids'].to(device)
attention_mask = batch['attention_mask'].to(device)
labels = batch['label'].to(device)
optimizer.zero_grad()
logits = model(input_ids, attention_mask)
loss = criterion(logits, labels)
loss.backward()
optimizer.step()
scheduler.step()
preds = torch.argmax(logits, dim=1).detach().cpu().numpy()
all_preds.extend(preds)
all_labels.extend(labels.detach().cpu().numpy())
running_loss += loss.item()
if i % 10 == 9:
avg_loss = running_loss / 10
acc = accuracy_score(all_labels[-100:], all_preds[-100:])
logging.info(f"Epoch [{epoch}] Batch [{i}] Loss: {avg_loss:.4f}, Acc: {acc:.4f}")
writer.add_scalar('Loss/train', avg_loss, epoch * len(dataloader) + i)
writer.add_scalar('Acc/train', acc, epoch * len(dataloader) + i)
running_loss = 0.0
all_preds, all_labels = [], []
@torch.no_grad()
def validate(model, dataloader, criterion, device, epoch, writer):
model.eval()
total_loss = 0.0
all_preds, all_labels = [], []
for batch in dataloader:
input_ids = batch['input_ids'].to(device)
attention_mask = batch['attention_mask'].to(device)
labels = batch['label'].to(device)
logits = model(input_ids, attention_mask)
loss = criterion(logits, labels)
total_loss += loss.item()
preds = torch.argmax(logits, dim=1).cpu().numpy()
all_preds.extend(preds)
all_labels.extend(labels.cpu().numpy())
avg_loss = total_loss / len(dataloader)
acc = accuracy_score(all_labels, all_preds)
logging.info(f"Epoch [{epoch}] Validation Loss: {avg_loss:.4f}, Val Acc: {acc:.4f}")
writer.add_scalar('Loss/val', avg_loss, epoch)
writer.add_scalar('Acc/val', acc, epoch)
return avg_loss, acc
# ================== Step 6: 主程序入口 ==================
def main():
setup_logger(config.log_dir)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
logging.info(f"Using device: {device}")
train_loader, val_loader, tokenizer = get_dataloaders(config.batch_size)
logging.info("Initializing BERT model...")
model = BERTClassifier(config.model_name, config.num_classes).to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.AdamW(model.parameters(), lr=config.lr, weight_decay=config.weight_decay)
scheduler = optim.lr_scheduler.LinearLR(optimizer, total_iters=len(train_loader) * config.epochs)
writer = SummaryWriter(log_dir=config.log_dir)
best_val_acc = 0.0
os.makedirs(config.checkpoint_dir, exist_ok=True)
logging.info("Start training...")
for epoch in range(1, config.epochs + 1):
logging.info(f"\nEpoch [{epoch}/{config.epochs}]")
train_one_epoch(model, train_loader, criterion, optimizer, scheduler, device, epoch, writer)
val_loss, val_acc = validate(model, val_loader, criterion, device, epoch, writer)
if val_acc > best_val_acc:
best_val_acc = val_acc
save_checkpoint({
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
'scheduler_state_dict': scheduler.state_dict(),
'epoch': epoch,
'val_acc': val_acc,
'val_loss': val_loss
}, config.checkpoint_dir)
logging.info(f"Saved best model with Val Acc: {best_val_acc:.4f}")
writer.close()
logging.info("Training complete.")
if __name__ == "__main__":
main()关键说明
| 模块 | 功能 |
|---|---|
| Config 类 | 集中式配置管理 |
| TextDataset + collate_fn | 支持任意格式的文本数据 |
| BERTClassifier | 使用 HuggingFace 的 BERT 模型作为基础架构 |
| train_one_epoch() | 模块化的训练逻辑 |
| validate() | 验证逻辑，输出准确率 |
| logging + SummaryWriter | 输出日志和 TensorBoard 可视化 |
| save_checkpoint() | 保存最佳模型状态 |
确保安装以下依赖：
pip install torch transformers scikit-learn pandas tensorboard启动 TensorBoard 查看训练过程：
tensorboard --logdir runs扩展建议
| 功能 | 描述 |
|---|---|
| 更复杂的任务 | 替换为命名实体识别（NER）、问答（QA）、摘要生成 |
| 更强的 Tokenizer | 使用 AutoTokenizer、RobertaTokenizer 等 |
| 混合精度训练 | 使用 torch.cuda.amp 提升训练效率 |
| 多 GPU 支持 | 使用 DataParallel 或 DistributedDataParallel |
| 推理模式 | 添加 predict() 函数用于推理阶段 |
| 多任务学习 | 构建统一输入接口，支持多个下游任务 |
总结
| 核心要素 | 在本项目中的体现 |
|---|---|
| 数据抽象 | TextDataset 可适配任意文本分类数据集 |
| 模型抽象 | BERTClassifier 可替换为其他预训练模型 |
| 损失统一 | 使用标准 CrossEntropyLoss |
| 日志统一 | Logging + TensorBoard |
| 模型保存 | Checkpoint 机制完善 |
| 可扩展性强 | 支持 NER、QA、摘要、多任务等 |
AutoRec
AutoRec（Autoencoders for Collaborative Filtering）是一种用于推荐系统的自动编码器模型，主要用于评分预测。它通过用户对项目的评分矩阵来学习用户的偏好或项目的特征。
主要改动点
- 数据集：需要一个用户-项目评分矩阵
- 模型结构：AutoRec 模型包含一个编码器和解码器，输入是评分向量，输出也是评分向量
- 损失函数：通常使用均方误差（MSE）作为损失函数
- 评估指标：除了准确率之外，还可以考虑 RMSE（均方根误差）等指标
完整代码示例
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from torch.utils.tensorboard import SummaryWriter
import numpy as np
import logging
import os
import time
# ================== Step 1: 自定义数据集 ==================
class RatingDataset(Dataset):
def __init__(self, ratings):
self.ratings = ratings
def __len__(self):
return len(self.ratings)
def __getitem__(self, idx):
user_ratings = self.ratings[idx]
# 将缺失值设为0
mask = np.not_equal(user_ratings, 0).astype(float)
return torch.FloatTensor(user_ratings), torch.FloatTensor(mask)
# ================== Step 2: 构建模型 ==================
class AutoRec(nn.Module):
def __init__(self, num_hidden=500, num_users=3706):
super(AutoRec, self).__init__()
self.encoder = nn.Linear(num_users, num_hidden)
self.decoder = nn.Linear(num_hidden, num_users)
def forward(self, x):
x = torch.sigmoid(self.encoder(x))
x = self.decoder(x)
return x
# 其他部分与原始模板相同，但需要做适当修改
# ================== 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 train_one_epoch(model, dataloader, criterion, optimizer, device, epoch, writer):
model.train()
running_loss = 0.0
for i, (inputs, masks) in enumerate(dataloader):
inputs, masks = inputs.to(device), masks.to(device)
optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs * masks, inputs * masks)
loss.backward()
optimizer.step()
running_loss += loss.item()
if i % 10 == 9:
avg_loss = running_loss / 10
logging.info(f"Epoch [{epoch}] Batch [{i}] Loss: {avg_loss:.4f}")
writer.add_scalar('Loss/train', avg_loss, epoch * len(dataloader) + i)
running_loss = 0.0
@torch.no_grad()
def validate(model, dataloader, criterion, device, epoch, writer):
model.eval()
val_loss = 0.0
for inputs, masks in dataloader:
inputs, masks = inputs.to(device), masks.to(device)
outputs = model(inputs)
loss = criterion(outputs * masks, inputs * masks)
val_loss += loss.item()
avg_loss = val_loss / len(dataloader)
logging.info(f"Epoch [{epoch}] Validation Loss: {avg_loss:.4f}")
writer.add_scalar('Loss/val', avg_loss, epoch)
return avg_loss
# ================== Step 5: 主程序入口 ==================
def main():
setup_logger()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
logging.info(f"Using device: {device}")
batch_size = 32
num_epochs = 10
lr = 0.001
# 假设我们有一个 6040 用户的评分数据
ratings = np.random.randint(0, 6, size=(6040, 3706)) # 示例数据
train_loader = DataLoader(RatingDataset(ratings), batch_size=batch_size, shuffle=True)
model = AutoRec().to(device)
criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=lr)
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.1)
writer = SummaryWriter(log_dir="runs")
best_loss = float('inf')
save_dir = "checkpoints"
os.makedirs(save_dir, exist_ok=True)
for epoch in range(num_epochs):
logging.info(f"\nEpoch [{epoch+1}/{num_epochs}]")
train_one_epoch(model, train_loader, criterion, optimizer, device, epoch+1, writer)
val_loss = validate(model, train_loader, criterion, device, epoch+1, writer)
scheduler.step()
if val_loss < best_loss:
best_loss = val_loss
torch.save(model.state_dict(), os.path.join(save_dir, f"best_model.pth"))
logging.info(f"Saved best model with loss: {best_loss:.4f}")
writer.close()
logging.info("Training complete.")
if __name__ == "__main__":
main()关键说明
| 模块 | 功能 |
|---|---|
| RatingDataset | 加载用户评分数据，通过 mask 处理缺失值 |
| AutoRec | 简单的编码器-解码器结构，用于评分预测 |
| 损失函数 | MSELoss，配合 mask 仅计算已有评分的损失 |
| logging + SummaryWriter | 同时输出到终端和 TensorBoard |
确保安装以下依赖：
pip install torch numpy tensorboard启动 TensorBoard 查看训练过程：
tensorboard --logdir runs扩展建议
| 功能 | 描述 |
|---|---|
| 更复杂的网络 | 增加隐藏层数或使用更深层的编码器-解码器 |
| 正则化 | 添加 L2 正则化或 Dropout 防止过拟合 |
| 评估指标 | 添加 RMSE、MAE 等推荐系统常用指标 |
| 真实数据 | 替换为 MovieLens 等真实评分数据集 |
| 物品视角 | 实现 Item-based AutoRec |
Neural Collaborative Filtering (NCF)
NCF（Neural Collaborative Filtering）将协同过滤与神经网络结合，使用 GMF（广义矩阵分解）和 MLP（多层感知机）建模用户与物品的交互，预测用户对某个物品的兴趣概率（点击/评分等）。
- 年份：2017
- 提出者：Xiangnan He 等（新加坡国立大学）
模型结构简述
NCF 的核心结构包括两个子网络：
- GMF (Generalized Matrix Factorization)
- 用户向量
user_emb和物品向量item_emb进行逐元素乘积 - 输出一个低维交互向量
- 用户向量
- MLP (Multi-Layer Perceptron)
- 用户和物品嵌入拼接后输入到全连接层中
- 提取高阶非线性交互特征
- 融合层 (Concat + FC)
- 将 GMF 和 MLP 的输出拼接，通过全连接层进行融合和预测
完整代码示例
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from torch.utils.tensorboard import SummaryWriter
import numpy as np
import logging
import os
import time
# ================== Step 1: 自定义数据集 ==================
class NCFTensorDataset(Dataset):
def __init__(self, user_ids, item_ids, labels, num_users, num_items):
self.user_ids = user_ids # shape: [num_samples]
self.item_ids = item_ids
self.labels = labels # shape: [num_samples]
self.num_users = num_users
self.num_items = num_items
def __len__(self):
return len(self.labels)
def __getitem__(self, idx):
user_id = self.user_ids[idx]
item_id = self.item_ids[idx]
label = self.labels[idx]
return user_id, item_id, label
# ================== Step 2: 构建模型 ==================
class NeuralCollaborativeFiltering(nn.Module):
def __init__(self, num_users, num_items, embedding_dim=8):
super(NeuralCollaborativeFiltering, self).__init__()
# Embedding 层
self.user_embedding = nn.Embedding(num_users, embedding_dim)
self.item_embedding = nn.Embedding(num_items, embedding_dim)
# GMF 分支
self.gmf_output = nn.Linear(embedding_dim, 1)
# MLP 分支
self.mlp = nn.Sequential(
nn.Linear(2 * embedding_dim, 64),
nn.ReLU(),
nn.Linear(64, 32),
nn.ReLU(),
nn.Linear(32, 16),
nn.ReLU()
)
# 融合层
self.final = nn.Linear(embedding_dim + 16, 1)
def forward(self, user_ids, item_ids):
# 获取嵌入
user_emb = self.user_embedding(user_ids) # [batch_size, embedding_dim]
item_emb = self.item_embedding(item_ids) # [batch_size, embedding_dim]
# GMF 部分
gmf_vector = user_emb * item_emb # element-wise product
gmf_out = self.gmf_output(gmf_vector).squeeze(-1) # [batch_size]
# MLP 部分
mlp_input = torch.cat([user_emb, item_emb], dim=-1) # [batch_size, 2*embedding_dim]
mlp_out = self.mlp(mlp_input) # [batch_size, 16]
# 融合
combined = torch.cat([gmf_vector, mlp_out], dim=-1) # [batch_size, embedding_dim + 16]
final_out = self.final(combined).squeeze(-1) # [batch_size]
return final_out
# ================== Step 3: 工具函数与配置 ==================
def get_dataloaders(batch_size=32):
# 示例数据：假设我们有 1000 条用户-物品交互记录，共 100 用户、50 物品
num_users = 100
num_items = 50
num_samples = 1000
user_ids = torch.randint(0, num_users, (num_samples,))
item_ids = torch.randint(0, num_items, (num_samples,))
labels = torch.rand(num_samples) # 假设是评分或点击概率，也可以用 0/1 表示正负样本
dataset = NCFTensorDataset(user_ids, item_ids, labels, num_users, num_items)
train_loader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
return train_loader, num_users, num_items
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 train_one_epoch(model, dataloader, criterion, optimizer, device, epoch, writer):
model.train()
running_loss = 0.0
for i, (user_ids, item_ids, labels) in enumerate(dataloader):
user_ids, item_ids, labels = user_ids.to(device), item_ids.to(device), labels.to(device)
optimizer.zero_grad()
outputs = model(user_ids, item_ids)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
running_loss += loss.item()
if i % 10 == 9:
avg_loss = running_loss / 10
logging.info(f"Epoch [{epoch}] Batch [{i}] Loss: {avg_loss:.4f}")
writer.add_scalar('Loss/train', avg_loss, epoch * len(dataloader) + i)
running_loss = 0.0
@torch.no_grad()
def validate(model, dataloader, criterion, device, epoch, writer):
model.eval()
val_loss = 0.0
for user_ids, item_ids, labels in dataloader:
user_ids, item_ids, labels = user_ids.to(device), item_ids.to(device), labels.to(device)
outputs = model(user_ids, item_ids)
loss = criterion(outputs, labels)
val_loss += loss.item()
avg_loss = val_loss / len(dataloader)
logging.info(f"Epoch [{epoch}] Validation Loss: {avg_loss:.4f}")
writer.add_scalar('Loss/val', avg_loss, epoch)
return avg_loss
# ================== Step 5: 主程序入口 ==================
def main():
setup_logger()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
logging.info(f"Using device: {device}")
batch_size = 32
num_epochs = 10
lr = 0.001
embedding_dim = 8
train_loader, num_users, num_items = get_dataloaders(batch_size)
model = NeuralCollaborativeFiltering(num_users, num_items, embedding_dim).to(device)
criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=lr)
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.1)
writer = SummaryWriter(log_dir="runs")
best_loss = float('inf')
save_dir = "checkpoints"
os.makedirs(save_dir, exist_ok=True)
for epoch in range(num_epochs):
logging.info(f"\nEpoch [{epoch+1}/{num_epochs}]")
train_one_epoch(model, train_loader, criterion, optimizer, device, epoch+1, writer)
val_loss = validate(model, train_loader, criterion, device, epoch+1, writer)
scheduler.step()
if val_loss < best_loss:
best_loss = val_loss
torch.save(model.state_dict(), os.path.join(save_dir, f"best_model.pth"))
logging.info(f"Saved best model with validation loss: {best_loss:.4f}")
writer.close()
logging.info("Training complete.")
if __name__ == "__main__":
main()关键说明
| 模块 | 功能 |
|---|---|
| NCFTensorDataset | 以用户 ID、物品 ID、标签三元组方式组织数据 |
| GMF 分支 | 用户和物品 Embedding 逐元素乘积，提取低阶交互 |
| MLP 分支 | 用户和物品 Embedding 拼接后经多层感知机提取高阶交互 |
| 融合层 | 拼接 GMF 向量和 MLP 输出，全连接层融合预测 |
| logging + SummaryWriter | 输出日志和 TensorBoard 可视化 |
确保安装以下依赖：
pip install torch numpy tensorboard启动 TensorBoard 查看训练过程：
tensorboard --logdir runs扩展建议
| 功能 | 描述 |
|---|---|
| 负采样 | 引入负样本提高模型区分能力 |
| 二分类 | 将 MSE 替换为 BCEWithLogitsLoss 做点击率预估 |
| 预训练初始化 | 用 GMF 和 MLP 分别预训练后融合 |
| 真实数据 | 替换为 MovieLens-1M 等真实推荐数据集 |
| 评估指标 | 添加 HR@K、NDCG@K 等推荐系统指标 |
Wide & Deep
Wide & Deep 模型是 Google 提出的推荐系统架构，将 Wide 部分（线性模型，记忆能力）和 Deep 部分（深度神经网络，泛化能力）结合，用于评分预测或点击率预估。
- 年份：2016
- 提出者：Google（用于 Google Play 应用推荐）
模型结构简述
| 组件 | 描述 |
|---|---|
| Wide 部分 | 线性层，输入是稠密特征或 one-hot/Embedding 后的交叉特征 |
| Deep 部分 | 多层感知机，输入是连续特征或 Embedding 特征 |
| 融合层 | 将 Wide 和 Deep 的输出拼接后输入到一个全连接层中 |
完整代码示例
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from torch.utils.tensorboard import SummaryWriter
import numpy as np
import logging
import os
import time
# ================== Step 1: 自定义数据集 ==================
class WideDeepDataset(Dataset):
def __init__(self, categorical_data, numerical_data, labels):
self.categorical = torch.LongTensor(categorical_data) # 类别型特征 [batch_size, num_categorical]
self.numerical = torch.FloatTensor(numerical_data) # 数值型特征 [batch_size, num_numerical]
self.labels = torch.FloatTensor(labels) # 标签（评分或点击率）[batch_size]
def __len__(self):
return len(self.labels)
def __getitem__(self, idx):
return (
self.categorical[idx],
self.numerical[idx],
self.labels[idx]
)
# ================== Step 2: 构建模型 ==================
class WideAndDeepModel(nn.Module):
def __init__(self, num_categories_list, embedding_dim=8, numerical_dim=5, hidden_dims=[128, 64], output_dim=1):
"""
:param num_categories_list: list of int, 每个类别特征的取值数量
:param embedding_dim: 每个类别特征嵌入维度
:param numerical_dim: 数值特征维度
:param hidden_dims: Deep部分隐藏层维度
"""
super(WideAndDeepModel, self).__init__()
# Embedding Layers for Categorical Features
self.embeddings = nn.ModuleList([
nn.Embedding(num_embeddings=dim, embedding_dim=embedding_dim)
for dim in num_categories_list
])
# Wide Layer (Linear)
wide_input_dim = sum(num_categories_list) + numerical_dim
self.wide = nn.Linear(wide_input_dim, 1)
# Deep Layers
deep_input_dim = len(num_categories_list) * embedding_dim + numerical_dim
layers = []
for h_dim in hidden_dims:
layers.append(nn.Linear(deep_input_dim, h_dim))
layers.append(nn.ReLU())
deep_input_dim = h_dim
self.deep = nn.Sequential(*layers)
# Final Fusion Layer
self.final = nn.Linear(hidden_dims[-1] + 1, output_dim)
def forward(self, categorical_x, numerical_x):
# Wide Part: One-hot + Concat + Linear
wide_inputs = [
torch.nn.functional.one_hot(torch.clamp(x, min=0), num_classes=num_cat).float()
for x, num_cat in zip(categorical_x.T, [emb.weight.shape[0] for emb in self.embeddings])]
wide_input = torch.cat(wide_inputs + [numerical_x], dim=1)
wide_out = self.wide(wide_input)
# Deep Part: Embedding + Concat + MLP
embeds = [emb(categorical_x[:, i]) for i, emb in enumerate(self.embeddings)]
deep_input = torch.cat(embeds + [numerical_x], dim=1)
deep_out = self.deep(deep_input)
# Fusion
combined = torch.cat([wide_out, deep_out], dim=1)
output = self.final(combined).squeeze(-1)
return output
# ================== Step 3: 工具函数与配置 ==================
def get_dataloaders(batch_size=32):
# 示例数据：假设有 3 个类别特征，每个类别分别有 100、50、20 种取值；2 个数值特征
num_samples = 1000
num_categories_list = [100, 50, 20]
num_numerical = 2
categorical_data = np.stack([
np.random.randint(0, n, size=num_samples)
for n in num_categories_list
], axis=1)
numerical_data = np.random.rand(num_samples, num_numerical)
labels = np.random.rand(num_samples) # 假设为评分任务
dataset = WideDeepDataset(categorical_data, numerical_data, labels)
train_loader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
return train_loader, num_categories_list, num_numerical
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 train_one_epoch(model, dataloader, criterion, optimizer, device, epoch, writer):
model.train()
running_loss = 0.0
for i, (categorical_x, numerical_x, labels) in enumerate(dataloader):
categorical_x, numerical_x, labels = (
categorical_x.to(device),
numerical_x.to(device),
labels.to(device)
)
optimizer.zero_grad()
outputs = model(categorical_x, numerical_x)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
running_loss += loss.item()
if i % 10 == 9:
avg_loss = running_loss / 10
logging.info(f"Epoch [{epoch}] Batch [{i}] Loss: {avg_loss:.4f}")
writer.add_scalar('Loss/train', avg_loss, epoch * len(dataloader) + i)
running_loss = 0.0
@torch.no_grad()
def validate(model, dataloader, criterion, device, epoch, writer):
model.eval()
val_loss = 0.0
for categorical_x, numerical_x, labels in dataloader:
categorical_x, numerical_x, labels = (
categorical_x.to(device),
numerical_x.to(device),
labels.to(device)
)
outputs = model(categorical_x, numerical_x)
loss = criterion(outputs, labels)
val_loss += loss.item()
avg_loss = val_loss / len(dataloader)
logging.info(f"Epoch [{epoch}] Validation Loss: {avg_loss:.4f}")
writer.add_scalar('Loss/val', avg_loss, epoch)
return avg_loss
# ================== Step 5: 主程序入口 ==================
def main():
setup_logger()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
logging.info(f"Using device: {device}")
batch_size = 32
num_epochs = 10
lr = 0.001
train_loader, num_categories_list, num_numerical = get_dataloaders(batch_size)
model = WideAndDeepModel(num_categories_list, numerical_dim=num_numerical).to(device)
criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=lr)
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.1)
writer = SummaryWriter(log_dir="runs")
best_loss = float('inf')
save_dir = "checkpoints"
os.makedirs(save_dir, exist_ok=True)
for epoch in range(num_epochs):
logging.info(f"\nEpoch [{epoch+1}/{num_epochs}]")
train_one_epoch(model, train_loader, criterion, optimizer, device, epoch+1, writer)
val_loss = validate(model, train_loader, criterion, device, epoch+1, writer)
scheduler.step()
if val_loss < best_loss:
best_loss = val_loss
torch.save(model.state_dict(), os.path.join(save_dir, f"best_model.pth"))
logging.info(f"Saved best model with validation loss: {best_loss:.4f}")
writer.close()
logging.info("Training complete.")
if __name__ == "__main__":
main()关键说明
| 模块 | 功能 |
|---|---|
| WideDeepDataset | 同时处理类别型特征和数值型特征 |
| Wide 部分 | One-hot + 线性层，记忆共现特征 |
| Deep 部分 | Embedding + MLP，泛化潜在特征关系 |
| 融合层 | 拼接 Wide 和 Deep 输出后进行预测 |
| logging + SummaryWriter | 输出日志和 TensorBoard 可视化 |
确保安装以下依赖：
pip install torch numpy tensorboard启动 TensorBoard 查看训练过程：
tensorboard --logdir runs扩展建议
| 功能 | 描述 |
|---|---|
| 交叉特征 | 手工构造特征交叉增强 Wide 部分 |
| 二分类 | 替换为 BCEWithLogitsLoss 做 CTR 预估 |
| 特征工程 | 引入更多连续特征和类别特征 |
| 真实数据 | 替换为 Criteo 或 Avazu 等推荐数据集 |
| 负采样 | 在点击率预估中引入正负样本平衡 |
DeepFM
DeepFM 模型将因子分解机（FM）和深度神经网络（DNN）结合，FM 部分建模低阶特征交互，DNN 部分建模高阶非线性特征交互，且不需要人工构造交叉特征。
- 年份：2017
- 提出者：Huawei Noah’s Ark Lab
模型结构简述
| 组件 | 描述 |
|---|---|
| FM 部分 | 使用 Embedding 向量之间的内积建模二阶特征交叉 |
| Deep 部分 | 多层感知机，输入是拼接的 Embedding 向量 |
| 融合层 | 将 FM 和 Deep 的输出相加并进行最终预测 |
与 Wide & Deep 不同的是，DeepFM 不需要人工构造交叉特征，而是通过 Embedding 自动建模特征组合。
完整代码示例
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from torch.utils.tensorboard import SummaryWriter
import numpy as np
import logging
import os
import time
# ================== Step 1: 自定义数据集 ==================
class DeepFMDataset(Dataset):
def __init__(self, categorical_data, numerical_data, labels):
"""
:param categorical_data: [num_samples, num_categorical]
:param numerical_data: [num_samples, num_numerical]
:param labels: [num_samples]
"""
self.categorical = torch.LongTensor(categorical_data)
self.numerical = torch.FloatTensor(numerical_data)
self.labels = torch.FloatTensor(labels)
def __len__(self):
return len(self.labels)
def __getitem__(self, idx):
return (
self.categorical[idx],
self.numerical[idx],
self.labels[idx]
)
# ================== Step 2: 构建模型 ==================
class DeepFM(nn.Module):
def __init__(self, num_categories_list, embedding_dim=8,
numerical_dim=5, hidden_dims=[128, 64], dropout=0.2):
"""
:param num_categories_list: list of int, 每个类别特征的取值数量
:param embedding_dim: 每个类别特征嵌入维度
:param numerical_dim: 数值特征维度
:param hidden_dims: Deep部分隐藏层维度
"""
super(DeepFM, self).__init__()
self.num_categorical = len(num_categories_list)
self.embedding_dim = embedding_dim
# Embedding Layers for Categorical Features
self.embeddings = nn.ModuleList([
nn.Embedding(num_embeddings=dim, embedding_dim=embedding_dim)
for dim in num_categories_list
])
# FM Layer: First-order + Second-order term
self.first_order_embeddings = nn.ModuleList([
nn.Embedding(num_embeddings=dim, embedding_dim=1)
for dim in num_categories_list
])
self.bias = nn.Parameter(torch.zeros(1))
# Deep Part
deep_input_dim = (self.num_categorical * embedding_dim) + numerical_dim
layers = []
for h_dim in hidden_dims:
layers.append(nn.Linear(deep_input_dim, h_dim))
layers.append(nn.BatchNorm1d(h_dim))
layers.append(nn.ReLU())
layers.append(nn.Dropout(dropout))
deep_input_dim = h_dim
self.deep = nn.Sequential(*layers)
# Final Fusion Layer
self.final = nn.Linear(hidden_dims[-1] + 1, 1)
def forward(self, categorical_x, numerical_x):
"""
:param categorical_x: [batch_size, num_categorical]
:param numerical_x: [batch_size, num_numerical]
"""
# ============ FM Part ============
# First-order term
first_order = [
self.first_order_embeddings[i](categorical_x[:, i])
for i in range(self.num_categorical)
]
first_order = torch.cat(first_order, dim=1).sum(dim=1)
# Second-order term
embeds = torch.cat([self.embeddings[i](categorical_x[:, i]).unsqueeze(1)
for i in range(self.num_categorical)], dim=1)
sum_of_squares = embeds.sum(dim=1).pow(2)
square_of_sums = embeds.pow(2).sum(dim=1)
second_order = 0.5 * (sum_of_squares - square_of_sums).sum(dim=1)
# ============ Deep Part ============
embeds_flattened = embeds.view(embeds.shape[0], -1)
deep_input = torch.cat([embeds_flattened, numerical_x], dim=1)
deep_out = self.deep(deep_input)
# ============ Fusion ============
combined = torch.cat([first_order.unsqueeze(1) + second_order.unsqueeze(1), deep_out], dim=1)
output = self.final(combined).squeeze(-1)
return output
# ================== Step 3: 工具函数与配置 ==================
def get_dataloaders(batch_size=32):
# 示例数据：假设有 3 个类别特征，每个类别分别有 100、50、20 种取值；2 个数值特征
num_samples = 1000
num_categories_list = [100, 50, 20]
num_numerical = 2
categorical_data = np.stack([
np.random.randint(0, n, size=num_samples)
for n in num_categories_list
], axis=1)
numerical_data = np.random.rand(num_samples, num_numerical)
labels = np.random.rand(num_samples) # 假设为评分任务
dataset = DeepFMDataset(categorical_data, numerical_data, labels)
train_loader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
return train_loader, num_categories_list, num_numerical
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 train_one_epoch(model, dataloader, criterion, optimizer, device, epoch, writer):
model.train()
running_loss = 0.0
for i, (categorical_x, numerical_x, labels) in enumerate(dataloader):
categorical_x, numerical_x, labels = (
categorical_x.to(device),
numerical_x.to(device),
labels.to(device)
)
optimizer.zero_grad()
outputs = model(categorical_x, numerical_x)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
running_loss += loss.item()
if i % 10 == 9:
avg_loss = running_loss / 10
logging.info(f"Epoch [{epoch}] Batch [{i}] Loss: {avg_loss:.4f}")
writer.add_scalar('Loss/train', avg_loss, epoch * len(dataloader) + i)
running_loss = 0.0
@torch.no_grad()
def validate(model, dataloader, criterion, device, epoch, writer):
model.eval()
val_loss = 0.0
for categorical_x, numerical_x, labels in dataloader:
categorical_x, numerical_x, labels = (
categorical_x.to(device),
numerical_x.to(device),
labels.to(device)
)
outputs = model(categorical_x, numerical_x)
loss = criterion(outputs, labels)
val_loss += loss.item()
avg_loss = val_loss / len(dataloader)
logging.info(f"Epoch [{epoch}] Validation Loss: {avg_loss:.4f}")
writer.add_scalar('Loss/val', avg_loss, epoch)
return avg_loss
# ================== Step 5: 主程序入口 ==================
def main():
setup_logger()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
logging.info(f"Using device: {device}")
batch_size = 32
num_epochs = 10
lr = 0.001
embedding_dim = 8
train_loader, num_categories_list, num_numerical = get_dataloaders(batch_size)
model = DeepFM(num_categories_list, numerical_dim=num_numerical, embedding_dim=embedding_dim).to(device)
criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=lr)
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.1)
writer = SummaryWriter(log_dir="runs")
best_loss = float('inf')
save_dir = "checkpoints"
os.makedirs(save_dir, exist_ok=True)
for epoch in range(num_epochs):
logging.info(f"\nEpoch [{epoch+1}/{num_epochs}]")
train_one_epoch(model, train_loader, criterion, optimizer, device, epoch+1, writer)
val_loss = validate(model, train_loader, criterion, device, epoch+1, writer)
scheduler.step()
if val_loss < best_loss:
best_loss = val_loss
torch.save(model.state_dict(), os.path.join(save_dir, f"best_model.pth"))
logging.info(f"Saved best model with validation loss: {best_loss:.4f}")
writer.close()
logging.info("Training complete.")
if __name__ == "__main__":
main()关键说明
| 模块 | 功能 |
|---|---|
| DeepFMDataset | 同时处理类别型特征和数值型特征 |
| FM 一阶部分 | 每个类别特征单独 Embedding（维度 1），求和得到一阶线性项 |
| FM 二阶部分 | 通过 Embedding 向量内积自动建模特征交叉 |
| Deep 部分 | Embedding 拼接后经多层感知机（含 BN + Dropout）提取高阶特征 |
| 融合层 | 拼接 FM 输出和 Deep 输出进行最终预测 |
确保安装以下依赖：
pip install torch numpy tensorboard启动 TensorBoard 查看训练过程：
tensorboard --logdir runs扩展建议
| 功能 | 描述 |
|---|---|
| 二分类 | 替换为 BCEWithLogitsLoss 做 CTR 预估 |
| FTRL 优化器 | Wide 部分使用 FTRL、Deep 部分使用 Adam |
| 特征工程 | 增加更多特征字段和交叉特征 |
| 真实数据 | 替换为 Criteo 或 Avazu 等推荐数据集 |
| 模型调优 | 调整 Embedding 维度、隐藏层数量和神经元数 |
YouTube DNN Recommender
YouTube DNN（Deep Neural Networks for YouTube Recommendations）是 Google 提出的推荐系统经典架构，使用双塔结构（Two-Tower Model）分别建模用户和物品，通过内积计算相似度作为预测分数。
- 年份：2016
- 提出者：Google Research
模型结构简述
| 组件 | 描述 |
|---|---|
| 用户塔（User Tower） | 多层感知机（MLP），输入是用户特征 + 历史行为 |
| 物品塔（Item Tower） | 多层感知机（MLP），输入是物品特征 |
| 匹配方式 | 内积或点积（dot product）表示用户对物品的兴趣程度 |
YouTube DNN 是现代召回阶段（Recall）中广泛使用的模型基础，也被称为双塔模型（Two-Tower Model）。适用场景包括视频推荐、商品推荐、新闻推荐等大规模排序任务。
完整代码示例
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from torch.utils.tensorboard import SummaryWriter
import numpy as np
import logging
import os
import time
# ================== Step 1: 自定义数据集 ==================
class YouTubeDNN_Dataset(Dataset):
def __init__(self, user_features, item_features, labels):
"""
:param user_features: [num_samples, num_user_features]
:param item_features: [num_samples, num_item_features]
:param labels: [num_samples], 表示点击/偏好分数
"""
self.user_features = torch.FloatTensor(user_features)
self.item_features = torch.FloatTensor(item_features)
self.labels = torch.FloatTensor(labels)
def __len__(self):
return len(self.labels)
def __getitem__(self, idx):
return (
self.user_features[idx],
self.item_features[idx],
self.labels[idx]
)
# ================== Step 2: 构建模型 ==================
class TwoTowerModel(nn.Module):
def __init__(self, user_input_dim, item_input_dim, embedding_dim=64, hidden_dims=[128, 64]):
super(TwoTowerModel, self).__init__()
# 用户塔
self.user_tower = self._build_mlp(user_input_dim, embedding_dim, hidden_dims)
# 物品塔
self.item_tower = self._build_mlp(item_input_dim, embedding_dim, hidden_dims)
def _build_mlp(self, input_dim, output_dim, hidden_dims):
layers = []
for h_dim in hidden_dims:
layers.append(nn.Linear(input_dim, h_dim))
layers.append(nn.ReLU())
layers.append(nn.BatchNorm1d(h_dim))
input_dim = h_dim
layers.append(nn.Linear(input_dim, output_dim))
return nn.Sequential(*layers)
def forward(self, user_x, item_x):
user_emb = self.user_tower(user_x)
item_emb = self.item_tower(item_x)
logits = torch.sum(user_emb * item_emb, dim=-1) # 内积
return logits
def get_user_embedding(self, user_x):
return self.user_tower(user_x)
def get_item_embedding(self, item_x):
return self.item_tower(item_x)
# ================== Step 3: 工具函数与配置 ==================
def get_dataloaders(batch_size=32):
# 示例数据：用户特征维度为 10，物品特征维度为 8
num_samples = 1000
user_feature_dim = 10
item_feature_dim = 8
user_features = np.random.rand(num_samples, user_feature_dim)
item_features = np.random.rand(num_samples, item_feature_dim)
labels = np.random.randint(0, 2, size=num_samples).astype(np.float32) # 二分类（点击/未点击）
dataset = YouTubeDNN_Dataset(user_features, item_features, labels)
train_loader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
return train_loader, user_feature_dim, item_feature_dim
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 train_one_epoch(model, dataloader, criterion, optimizer, device, epoch, writer):
model.train()
running_loss = 0.0
for i, (user_x, item_x, labels) in enumerate(dataloader):
user_x, item_x, labels = user_x.to(device), item_x.to(device), labels.to(device)
optimizer.zero_grad()
outputs = model(user_x, item_x)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
running_loss += loss.item()
if i % 10 == 9:
avg_loss = running_loss / 10
logging.info(f"Epoch [{epoch}] Batch [{i}] Loss: {avg_loss:.4f}")
writer.add_scalar('Loss/train', avg_loss, epoch * len(dataloader) + i)
running_loss = 0.0
@torch.no_grad()
def validate(model, dataloader, criterion, device, epoch, writer):
model.eval()
val_loss = 0.0
for user_x, item_x, labels in dataloader:
user_x, item_x, labels = user_x.to(device), item_x.to(device), labels.to(device)
outputs = model(user_x, item_x)
loss = criterion(outputs, labels)
val_loss += loss.item()
avg_loss = val_loss / len(dataloader)
logging.info(f"Epoch [{epoch}] Validation Loss: {avg_loss:.4f}")
writer.add_scalar('Loss/val', avg_loss, epoch)
return avg_loss
# ================== Step 5: 主程序入口 ==================
def main():
setup_logger()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
logging.info(f"Using device: {device}")
batch_size = 32
num_epochs = 10
lr = 0.001
embedding_dim = 64
train_loader, user_dim, item_dim = get_dataloaders(batch_size)
model = TwoTowerModel(user_dim, item_dim, embedding_dim=embedding_dim).to(device)
criterion = nn.BCEWithLogitsLoss()
optimizer = optim.Adam(model.parameters(), lr=lr)
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.1)
writer = SummaryWriter(log_dir="runs")
best_loss = float('inf')
save_dir = "checkpoints"
os.makedirs(save_dir, exist_ok=True)
for epoch in range(num_epochs):
logging.info(f"\nEpoch [{epoch+1}/{num_epochs}]")
train_one_epoch(model, train_loader, criterion, optimizer, device, epoch+1, writer)
val_loss = validate(model, train_loader, criterion, device, epoch+1, writer)
scheduler.step()
if val_loss < best_loss:
best_loss = val_loss
torch.save(model.state_dict(), os.path.join(save_dir, f"best_model.pth"))
logging.info(f"Saved best model with validation loss: {best_loss:.4f}")
writer.close()
logging.info("Training complete.")
if __name__ == "__main__":
main()关键说明
| 模块 | 功能 |
|---|---|
| YouTubeDNN_Dataset | 用户特征 + 物品特征 + 标签三元组数据 |
| TwoTowerModel（用户塔） | 多层感知机将用户特征编码为向量 |
| TwoTowerModel（物品塔） | 多层感知机将物品特征编码为向量 |
| 内积匹配 | 用户向量与物品向量点积作为预测分数 |
| get_user_embedding / get_item_embedding | 提供离线复用，用于向量检索与召回 |
确保安装以下依赖：
pip install torch numpy tensorboard启动 TensorBoard 查看训练过程：
tensorboard --logdir runs扩展建议
| 功能 | 描述 |
|---|---|
| 召回阶段 | 将用户和物品向量输出到向量数据库（如 Faiss）进行召回 |
| 负采样 | 引入负样本提高模型区分能力 |
| 特征工程 | 加入更多用户侧和物品侧上下文特征 |
| 多任务学习 | 扩展为多目标学习，如同时预测点击和观看时长 |
| 真实数据 | 替换为 MovieLens 等真实推荐数据集 |
Graph Neural Networks (GNNs) for Recommendation
GNN 推荐系统将用户与物品之间的交互构建成图结构（如用户-物品二部图），利用图神经网络（如 GCN、GAT、LightGCN）聚合邻居信息，学习更丰富的节点嵌入，基于用户和物品的嵌入向量进行点击预测或评分预测。
代表模型包括 NGCF（Neural Graph Collaborative Filtering）、LightGCN、PinSAGE、GPR-GNN 以及 GAT-based 推荐模型。
模型结构简述（以 LightGCN 为例）
| 组件 | 描述 |
|---|---|
| 图结构构建 | 构造用户-物品邻接矩阵（edge_index） |
| GNN 层 | 使用 LightGCN 进行多层消息传递（无参数） |
| 融合 | 多层嵌入取平均作为最终表示 |
| 内积预测 | 用户向量与物品向量内积为点击概率 |
我们将使用 PyG（PyTorch Geometric）来构建图神经网络部分。
安装依赖
pip install torch torchvision torchaudio
pip install torch-geometric完整代码示例
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, Dataset
from torch_geometric.data import Data
from torch_geometric.nn import GCNConv, TopKPooling, global_mean_pool
from torch_geometric.nn import LightGCN
from torch_geometric.utils import train_test_split_edges
import numpy as np
import logging
import os
import time
from torch.utils.tensorboard import SummaryWriter
# ================== Step 1: 自定义数据集 ==================
class GNNDataset(Dataset):
def __init__(self, data):
self.data = data
def __len__(self):
return 1 # 整个图作为一个样本
def __getitem__(self, idx):
return self.data
# ================== Step 2: 构建模型 ==================
class GNNRecommender(nn.Module):
def __init__(self, num_users, num_items, embedding_dim=64, num_layers=3):
super(GNNRecommender, self).__init__()
self.model = LightGCN(num_nodes=num_users + num_items,
embedding_dim=embedding_dim,
num_layers=num_layers)
def forward(self, edge_index, edge_weight=None):
embeddings = self.model(edge_index, edge_weight)
return embeddings
def recommend(self, user_ids, top_k=10):
embeddings = self.model()
user_emb = embeddings[user_ids]
item_emb = embeddings[self.num_users:]
scores = torch.matmul(user_emb, item_emb.t())
_, top_indices = torch.topk(scores, k=top_k, dim=1)
return top_indices # 返回 top-k 物品 ID
@property
def num_users(self):
return self.model.num_nodes_list[0]
# ================== Step 3: 工具函数与配置 ==================
def build_graph_data(num_users=1000, num_items=5000, num_interactions=10000):
"""
构建用户-物品交互图
"""
user_ids = np.random.randint(0, num_users, size=num_interactions)
item_ids = np.random.randint(0, num_items, size=num_interactions)
edge_index = torch.tensor(np.stack([user_ids, item_ids + num_users], axis=0), dtype=torch.long)
edge_index = torch.cat([edge_index, edge_index.flip(0)], dim=1) # 双向边
data = Data(edge_index=edge_index)
data = train_test_split_edges(data, val_ratio=0.05, test_ratio=0.1)
data.num_users = num_users
data.num_items = num_items
data.total_num_nodes = num_users + num_items
return data
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 train_one_epoch(model, data, optimizer, device, epoch, writer):
model.train()
optimizer.zero_grad()
embeddings = model(data.edge_index.to(device))
user_emb = embeddings[:data.num_users]
item_emb = embeddings[data.num_users:data.num_users + data.num_items]
pos_edge_label_index = data.train_pos_edge_index.to(device)
neg_edge_label_index = torch.randint(0, data.num_items, (2, data.train_pos_edge_index.size(1)), device=device)
neg_edge_label_index[0] += data.num_users
pos_scores = (user_emb[pos_edge_label_index[0]] * item_emb[pos_edge_label_index[1] - data.num_users]).sum(dim=1)
neg_scores = (user_emb[neg_edge_label_index[0]] * item_emb[neg_edge_label_index[1] - data.num_users]).sum(dim=1)
loss = -torch.log(torch.sigmoid(pos_scores - neg_scores)).mean()
loss.backward()
optimizer.step()
logging.info(f"Epoch [{epoch}] Loss: {loss.item():.4f}")
writer.add_scalar('Loss/train', loss.item(), epoch)
return loss.item()
@torch.no_grad()
def validate(model, data, device, epoch, writer):
model.eval()
embeddings = model(data.edge_index.to(device))
user_emb = embeddings[:data.num_users]
item_emb = embeddings[data.num_users:data.num_users + data.num_items]
val_edge_label_index = data.val_pos_edge_index.to(device)
val_neg_edge_label_index = data.val_neg_edge_index.to(device)
pos_scores = (user_emb[val_edge_label_index[0]] * item_emb[val_edge_label_index[1] - data.num_users]).sum(dim=1)
neg_scores = (user_emb[val_neg_edge_label_index[0]] * item_emb[val_neg_edge_label_index[1] - data.num_users]).sum(dim=1)
acc = (pos_scores > neg_scores).float().mean().item()
logging.info(f"Epoch [{epoch}] Validation Accuracy: {acc:.4f}")
writer.add_scalar('Accuracy/val', acc, epoch)
return acc
# ================== Step 5: 主程序入口 ==================
def main():
setup_logger()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
logging.info(f"Using device: {device}")
num_users = 1000
num_items = 5000
embedding_dim = 64
num_epochs = 20
lr = 0.001
data = build_graph_data(num_users, num_items)
dataset = GNNDataset(data)
dataloader = DataLoader(dataset, batch_size=1, shuffle=False)
model = GNNRecommender(num_users, num_items, embedding_dim=embedding_dim, num_layers=3).to(device)
optimizer = optim.Adam(model.parameters(), lr=lr)
writer = SummaryWriter(log_dir="runs")
best_acc = 0.0
save_dir = "checkpoints"
os.makedirs(save_dir, exist_ok=True)
for epoch in range(num_epochs):
logging.info(f"\nEpoch [{epoch+1}/{num_epochs}]")
for data in dataloader:
train_one_epoch(model, data, optimizer, device, epoch+1, writer)
val_acc = validate(model, data, device, epoch+1, writer)
if val_acc > best_acc:
best_acc = val_acc
torch.save(model.state_dict(), os.path.join(save_dir, f"best_model.pth"))
logging.info(f"Saved best model with validation accuracy: {best_acc:.4f}")
writer.close()
logging.info("Training complete.")
if __name__ == "__main__":
main()关键说明
| 模块 | 功能 |
|---|---|
| build_graph_data | 构建用户-物品交互二部图，生成边索引 |
| LightGCN | 无参数图卷积，多层邻居信息聚合 |
| BPR Loss | 正负样本对内积差值，sigmoid + logloss |
| recommend() | 基于学习到的嵌入进行 Top-K 推荐 |
| train_test_split_edges | 自动划分训练/验证/测试边 |
扩展建议
| 功能 | 描述 |
|---|---|
| 异构图扩展 | 引入物品属性、用户画像等多种节点类型 |
| GAT/GCN | 将 LightGCN 替换为带有注意力机制的 GAT |
| 大规模图 | 使用 PinSage 或 GraphSAGE 进行子图采样 |
| 知识图谱 | 融合知识图谱增强物品表示 |
| 评估指标 | 添加 Recall@K、NDCG@K 等排序指标 |
RNN for Session-based Recommendation
Session-based 推荐系统给定用户在当前会话中点击/浏览的商品序列，预测下一个可能点击/购买的商品。数据由多个”会话”组成，每个会话是一个商品 ID 序列，没有明确的用户 ID，只依赖行为序列建模用户意图。
适用模型包括 RNN / GRU / LSTM、Self-Attention / Transformer 以及 GNN-based Session Recommender。
模型结构简述（以 GRU 为例）
| 组件 | 描述 |
|---|---|
| Embedding Layer | 将物品 ID 映射为向量 |
| RNN (GRU/LSTM) | 编码会话序列 |
| Output Layer | 预测下一个物品的概率分布 |
| Loss | CrossEntropyLoss 或 Sampled Softmax |
完整代码示例
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from torch.nn.utils.rnn import pad_sequence
import numpy as np
import logging
import os
import time
from torch.utils.tensorboard import SummaryWriter
# ================== Step 1: 自定义数据集 ==================
class SessionDataset(Dataset):
def __init__(self, sessions):
"""
:param sessions: list of lists, 每个元素是一个 session 商品 ID 序列
"""
self.sessions = sessions
def __len__(self):
return len(self.sessions)
def __getitem__(self, idx):
session = self.sessions[idx]
input_ids = session[:-1] # 输入序列 [x_1, x_2, ..., x_{n-1}]
target_id = session[-1] # 目标项 x_n
return torch.LongTensor(input_ids), target_id
def collate_fn(batch):
input_seqs, targets = zip(*batch)
input_padded = pad_sequence(input_seqs, batch_first=True)
input_lengths = torch.LongTensor([len(seq) for seq in input_seqs])
targets = torch.LongTensor(targets)
return input_padded, input_lengths, targets
# ================== Step 2: 构建模型 ==================
class RNNRecommender(nn.Module):
def __init__(self, num_items, embedding_dim=64, hidden_dim=128, num_layers=1, rnn_type='GRU'):
super(RNNRecommender, self).__init__()
self.embedding = nn.Embedding(num_embeddings=num_items + 1, embedding_dim=embedding_dim, padding_idx=0)
if rnn_type == 'GRU':
self.rnn = nn.GRU(input_size=embedding_dim, hidden_size=hidden_dim,
num_layers=num_layers, batch_first=True)
elif rnn_type == 'LSTM':
self.rnn = nn.LSTM(input_size=embedding_dim, hidden_size=hidden_dim,
num_layers=num_layers, batch_first=True)
else:
self.rnn = nn.RNN(input_size=embedding_dim, hidden_size=hidden_dim,
num_layers=num_layers, batch_first=True)
self.output = nn.Linear(hidden_dim, num_items)
def forward(self, input_seq, lengths):
embedded = self.embedding(input_seq)
packed = nn.utils.rnn.pack_padded_sequence(embedded, lengths, batch_first=True, enforce_sorted=False)
output, hidden = self.rnn(packed)
output, _ = nn.utils.rnn.pad_packed_sequence(output, batch_first=True)
if isinstance(hidden, tuple): # LSTM returns a tuple
hidden = hidden[0]
logits = self.output(hidden[-1]) # 取最后一个隐藏状态作为输出
return logits
# ================== Step 3: 工具函数与配置 ==================
def get_dataloader(batch_size=32, num_items=1000, session_length=5):
# 示例数据：随机生成 session 数据
num_sessions = 10000
sessions = []
for _ in range(num_sessions):
length = np.random.randint(2, session_length + 1)
session = np.random.randint(1, num_items + 1, size=length).tolist()
sessions.append(session)
dataset = SessionDataset(sessions)
dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True, collate_fn=collate_fn)
return dataloader, num_items
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 train_one_epoch(model, dataloader, criterion, optimizer, device, epoch, writer):
model.train()
running_loss = 0.0
for i, (input_seq, lengths, targets) in enumerate(dataloader):
input_seq, targets = input_seq.to(device), targets.to(device)
optimizer.zero_grad()
outputs = model(input_seq, lengths)
loss = criterion(outputs, targets)
loss.backward()
optimizer.step()
running_loss += loss.item()
if i % 10 == 9:
avg_loss = running_loss / 10
logging.info(f"Epoch [{epoch}] Batch [{i}] Loss: {avg_loss:.4f}")
writer.add_scalar('Loss/train', avg_loss, epoch * len(dataloader) + i)
running_loss = 0.0
@torch.no_grad()
def validate(model, dataloader, criterion, device, epoch, writer):
model.eval()
total_loss = 0.0
total_acc = 0.0
count = 0
for input_seq, lengths, targets in dataloader:
input_seq, targets = input_seq.to(device), targets.to(device)
outputs = model(input_seq, lengths)
loss = criterion(outputs, targets)
total_loss += loss.item() * len(targets)
preds = outputs.argmax(dim=-1)
acc = (preds == targets).float().mean().item()
total_acc += acc * len(targets)
count += len(targets)
avg_loss = total_loss / count
avg_acc = total_acc / count
logging.info(f"Epoch [{epoch}] Validation Loss: {avg_loss:.4f}, Accuracy: {avg_acc:.4f}")
writer.add_scalar('Loss/val', avg_loss, epoch)
writer.add_scalar('Accuracy/val', avg_acc, epoch)
return avg_loss, avg_acc
# ================== Step 5: 主程序入口 ==================
def main():
setup_logger()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
logging.info(f"Using device: {device}")
batch_size = 128
num_epochs = 20
lr = 0.001
embedding_dim = 64
hidden_dim = 128
num_items = 1000 # 假设有 1000 个不同的 item
rnn_type = 'GRU'
train_loader, num_items = get_dataloader(batch_size=batch_size, num_items=num_items)
model = RNNRecommender(num_items=num_items, embedding_dim=embedding_dim,
hidden_dim=hidden_dim, rnn_type=rnn_type).to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=lr)
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.5)
writer = SummaryWriter(log_dir="runs")
best_acc = 0.0
save_dir = "checkpoints"
os.makedirs(save_dir, exist_ok=True)
for epoch in range(num_epochs):
logging.info(f"\nEpoch [{epoch+1}/{num_epochs}]")
train_one_epoch(model, train_loader, criterion, optimizer, device, epoch+1, writer)
val_loss, val_acc = validate(model, train_loader, criterion, device, epoch+1, writer)
scheduler.step()
if val_acc > best_acc:
best_acc = val_acc
torch.save(model.state_dict(), os.path.join(save_dir, f"best_model.pth"))
logging.info(f"Saved best model with validation accuracy: {best_acc:.4f}")
writer.close()
logging.info("Training complete.")
if __name__ == "__main__":
main()关键说明
| 模块 | 功能 |
|---|---|
| SessionDataset | 会话序列数据，最后一项作为标签 |
| collate_fn | 变长序列填充（pad_sequence + pack_padded） |
| RNNRecommender | 支持 RNN/GRU/LSTM，取最后一个隐藏状态预测 |
| CrossEntropyLoss | 下一个物品的多分类损失 |
| 验证指标 | 准确率 + 损失 |
确保安装以下依赖：
pip install torch numpy tensorboard启动 TensorBoard 查看训练过程：
tensorboard --logdir runs扩展建议
| 功能 | 描述 |
|---|---|
| Attention 机制 | 在 RNN 输出上添加 Self-Attention 捕捉全局依赖 |
| Transformer | 替换 RNN 为 Transformer 编码器 |
| 负采样 | 使用 Sampled Softmax 减少计算量 |
| 真实数据 | 替换为 Yoochoose、Diginetica 等公开数据集 |
| 冷启动处理 | 融入物品 side information（如类别、价格） |
BERT4Rec
BERT4Rec（BERT for Sequential Recommendation）将推荐任务建模为序列到标签的预测问题，使用 Transformer 编码器结构进行双向建模（类似 BERT）。对用户行为序列中某个位置的商品进行掩码（Mask），然后通过上下文预测该商品。
- 提出者：Huawei Noah’s Ark Lab & University of Science and Technology of China
- 年份：2019
BERT4Rec 可以捕捉长距离依赖关系，利用双向上下文信息提升预测准确性，支持 Top-K 推荐、Next-item 预测等任务。
模型结构简述
| 组件 | 描述 |
|---|---|
| 输入嵌入 | Item ID + Positional Embedding |
| Transformer Encoder | 多层自注意力机制 |
| MLM Head | 预测被 Mask 的 item ID |
| 训练目标 | 使用 CrossEntropyLoss 预测 masked items |
完整代码示例
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from torch.nn.utils.rnn import pad_sequence
import numpy as np
import logging
import os
import time
from torch.utils.tensorboard import SummaryWriter
# ================== Step 1: 自定义数据集 ==================
class BERT4RecDataset(Dataset):
def __init__(self, sequences, max_len=50, mask_prob=0.15):
"""
:param sequences: list of lists, 用户行为序列列表
:param max_len: 序列最大长度
:param mask_prob: 掩码概率
"""
self.sequences = [s[-max_len:] for s in sequences]
self.max_len = max_len
self.mask_prob = mask_prob
self.num_items = max(max(seq) for seq in sequences) + 1 # 自动识别 item 数量
def __len__(self):
return len(self.sequences)
def __getitem__(self, idx):
seq = self.sequences[idx]
input_ids = torch.zeros(self.max_len, dtype=torch.long)
labels = torch.zeros(self.max_len, dtype=torch.long)
attention_mask = torch.zeros(self.max_len, dtype=torch.long)
input_ids[-len(seq):] = torch.LongTensor(seq)
attention_mask[-len(seq):] = 1
# 随机 Mask
labels = input_ids.clone()
indices = torch.rand(len(seq)) < self.mask_prob
input_ids[-len(seq):][indices] = self.num_items # 使用 [MASK] token
labels[-len(seq):][~indices] = -100 # 不计算未被 mask 的 loss
return {
"input_ids": input_ids,
"attention_mask": attention_mask,
"labels": labels
}
# ================== Step 2: 构建模型 ==================
class BERT4Rec(nn.Module):
def __init__(self, num_items, embedding_dim=64, max_len=50, num_heads=4, num_layers=3, dropout=0.1):
super(BERT4Rec, self).__init__()
self.item_embeddings = nn.Embedding(num_embeddings=num_items + 1, embedding_dim=embedding_dim, padding_idx=0)
self.position_embeddings = nn.Embedding(num_embeddings=max_len, embedding_dim=embedding_dim)
self.register_buffer("position_ids", torch.arange(max_len).expand((1, -1)))
self.encoder = nn.TransformerEncoder(
encoder_layer=nn.TransformerEncoderLayer(d_model=embedding_dim, nhead=num_heads, dropout=dropout),
num_layers=num_layers
)
self.ln = nn.LayerNorm(embedding_dim)
self.head = nn.Linear(embedding_dim, num_items)
def forward(self, input_ids, attention_mask=None):
# Embeddings
position_ids = self.position_ids[:, :input_ids.size(1)]
token_embeddings = self.item_embeddings(input_ids)
position_embeddings = self.position_embeddings(position_ids)
embeddings = token_embeddings + position_embeddings
if attention_mask is not None:
src_key_padding_mask = ~attention_mask.bool()
# Transformer 编码
outputs = self.encoder(
embeddings.transpose(0, 1), # [seq_len, batch_size, embed_dim]
src_key_padding_mask=src_key_padding_mask
).transpose(0, 1)
outputs = self.ln(outputs)
logits = self.head(outputs) # [batch_size, seq_len, num_items]
return logits
# ================== Step 3: 工具函数与配置 ==================
def get_dataloader(batch_size=32, max_len=50, mask_prob=0.15):
# 示例数据：随机生成用户行为序列
num_users = 10000
avg_seq_len = 20
num_items = 1000
sequences = []
for _ in range(num_users):
length = np.random.randint(5, avg_seq_len * 2)
seq = np.random.randint(1, num_items + 1, size=length).tolist()
sequences.append(seq)
dataset = BERT4RecDataset(sequences, max_len=max_len, mask_prob=mask_prob)
dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
return dataloader, num_items
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 train_one_epoch(model, dataloader, criterion, optimizer, device, epoch, writer):
model.train()
running_loss = 0.0
for i, batch in enumerate(dataloader):
input_ids = batch['input_ids'].to(device)
attention_mask = batch['attention_mask'].to(device)
labels = batch['labels'].to(device)
optimizer.zero_grad()
logits = model(input_ids, attention_mask)
loss = criterion(logits.view(-1, logits.size(-1)), labels.view(-1))
loss.backward()
optimizer.step()
running_loss += loss.item()
if i % 10 == 9:
avg_loss = running_loss / 10
logging.info(f"Epoch [{epoch}] Batch [{i}] Loss: {avg_loss:.4f}")
writer.add_scalar('Loss/train', avg_loss, epoch * len(dataloader) + i)
running_loss = 0.0
@torch.no_grad()
def validate(model, dataloader, criterion, device, epoch, writer):
model.eval()
total_loss = 0.0
total_acc = 0.0
count = 0
for batch in dataloader:
input_ids = batch['input_ids'].to(device)
attention_mask = batch['attention_mask'].to(device)
labels = batch['labels'].to(device)
logits = model(input_ids, attention_mask)
loss = criterion(logits.view(-1, logits.size(-1)), labels.view(-1))
total_loss += loss.item() * input_ids.size(0)
preds = logits.argmax(dim=-1)
valid_mask = (labels != -100)
acc = (preds[valid_mask] == labels[valid_mask]).float().mean().item()
total_acc += acc * input_ids.size(0)
count += input_ids.size(0)
avg_loss = total_loss / count
avg_acc = total_acc / count
logging.info(f"Epoch [{epoch}] Validation Loss: {avg_loss:.4f}, Accuracy: {avg_acc:.4f}")
writer.add_scalar('Loss/val', avg_loss, epoch)
writer.add_scalar('Accuracy/val', avg_acc, epoch)
return avg_loss, avg_acc
# ================== Step 5: 主程序入口 ==================
def main():
setup_logger()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
logging.info(f"Using device: {device}")
batch_size = 128
num_epochs = 20
lr = 0.001
embedding_dim = 64
max_len = 50
num_heads = 4
num_layers = 3
dropout = 0.1
train_loader, num_items = get_dataloader(batch_size=batch_size, max_len=max_len)
model = BERT4Rec(num_items=num_items, embedding_dim=embedding_dim,
max_len=max_len, num_heads=num_heads, num_layers=num_layers, dropout=dropout).to(device)
criterion = nn.CrossEntropyLoss(ignore_index=-100)
optimizer = optim.AdamW(model.parameters(), lr=lr)
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.5)
writer = SummaryWriter(log_dir="runs")
best_acc = 0.0
save_dir = "checkpoints"
os.makedirs(save_dir, exist_ok=True)
for epoch in range(num_epochs):
logging.info(f"\nEpoch [{epoch+1}/{num_epochs}]")
train_one_epoch(model, train_loader, criterion, optimizer, device, epoch+1, writer)
val_loss, val_acc = validate(model, train_loader, criterion, device, epoch+1, writer)
scheduler.step()
if val_acc > best_acc:
best_acc = val_acc
torch.save(model.state_dict(), os.path.join(save_dir, f"best_model.pth"))
logging.info(f"Saved best model with validation accuracy: {best_acc:.4f}")
writer.close()
logging.info("Training complete.")
if __name__ == "__main__":
main()关键说明
| 模块 | 功能 |
|---|---|
| BERT4RecDataset | 行为序列数据，随机 Mask 部分位置进行预测 |
| TransformerEncoder | 双向自注意力编码器，捕捉序列上下文 |
| Positional Embedding | 学习式位置编码，感知序列位置 |
| Cloze Task | 类比 BERT MLM，预测被 Mask 的商品 |
| AdamW 优化器 | 适用于 Transformer 架构的优化器 |
确保安装以下依赖：
pip install torch numpy tensorboard启动 TensorBoard 查看训练过程：
tensorboard --logdir runs扩展建议
| 功能 | 描述 |
|---|---|
| 时间信息 | 加入时间间隔编码增强序列建模 |
| 预训练 + 微调 | 在全局行为数据上预训练，再在特定任务微调 |
| SASRec | 替换为单向 Transformer 的自回归模型 |
| 真实数据 | 替换为 Amazon Reviews、MovieLens-1M 等序列数据集 |
| 评估指标 | 添加 HR@K、NDCG@K、MRR 等排序指标 |
Long Short-Term Memory（LSTM）
LSTM 是一种特殊的循环神经网络（RNN）结构，专门设计用来解决传统 RNN 在处理长序列时的梯度消失和梯度爆炸问题。
它能够记住长期依赖信息，在自然语言处理、语音识别、时间序列预测等任务中表现优异。
传统的 RNN 在理论上可以记忆前面的信息，但在实践中，当序列变长时，模型很难记住早期的信息。这是因为反向传播过程中，梯度会逐渐变小（梯度消失），导致模型无法学习到长期依赖关系。
LSTM 的核心思想：门控机制
LSTM 引入了”门控机制（Gating Mechanism）“，通过三个主要的”门”来控制信息的流动：
遗忘门（Forget Gate）
决定哪些信息从 cell state（记忆单元）中被丢弃。
- 使用 sigmoid 激活函数输出一个 0~1 的值：
0表示完全忘记1表示完全保留
输入门（Input Gate）
决定哪些新信息将被加入到 cell state 中。
包括两个部分：
- 一个 sigmoid 层决定哪些值要更新
- 一个 tanh 层创建一个新的候选值向量
输出门（Output Gate）
决定当前时刻的输出是什么。
- 首先用 sigmoid 层决定 cell state 的哪部分将被输出
- 然后将 cell state 通过 tanh 激活函数，并与 sigmoid 输出相乘，得到最终输出
模型结构
| 组件 | 说明 |
|---|---|
CustomLSTMCell | 自定义 LSTM 处理单元，实现完整的门控逻辑 |
CustomLSTM | 自定义 LSTM 模型层，支持多层堆叠 |
LSTMModel | 封装模型，取最后一个时间步输出进行分类 |
使用 PyTorch 实现 LSTM
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from torch.utils.tensorboard import SummaryWriter
import os
import logging
import time
# ================== Step 1: 自定义序列数据集 ==================
class SequenceDataset(Dataset):
def __init__(self, num_samples=1000, seq_len=20, input_dim=10, num_classes=3):
self.num_samples = num_samples
self.seq_len = seq_len
self.input_dim = input_dim
self.data = torch.randn(num_samples, seq_len, input_dim)
self.labels = torch.randint(0, num_classes, (num_samples,))
def __len__(self):
return self.num_samples
def __getitem__(self, idx):
return self.data[idx], self.labels[idx]
# ================== Step 2: 构建 LSTM 模型 ==================
class CustomLSTMCell(nn.Module):
'''
自定义LSTM处理单元
'''
def __init__(self, input_dim, hidden_dim):
super(CustomLSTMCell, self).__init__()
self.input_dim = input_dim
self.hidden_dim = hidden_dim
# 定义4个门的线性层：i (input), f (forget), g (candidate), o (output)
self.W_i = nn.Linear(input_dim + hidden_dim, hidden_dim)
self.W_f = nn.Linear(input_dim + hidden_dim, hidden_dim)
self.W_g = nn.Linear(input_dim + hidden_dim, hidden_dim)
self.W_o = nn.Linear(input_dim + hidden_dim, hidden_dim)
def forward(self, x, hidden):
h_prev, c_prev = hidden
# 拼接当前输入 x 和上一时刻隐藏状态 h_prev
combined = torch.cat((x, h_prev), dim=1)
# 遗忘门：决定丢弃多少之前的 cell state
f_t = torch.sigmoid(self.W_f(combined))
# 输入门：决定哪些新信息将被加入到 cell state 中
i_t = torch.sigmoid(self.W_i(combined))
g_t = torch.tanh(self.W_g(combined))
# 更新 cell state
c_t = f_t * c_prev + i_t * g_t
# 输出门：决定输出什么
o_t = torch.sigmoid(self.W_o(combined))
h_t = o_t * torch.tanh(c_t)
return h_t, c_t
class CustomLSTM(nn.Module):
'''
自定义LSTM模型层
'''
def __init__(self, input_dim, hidden_dim, num_layers=1, batch_first=True):
super(CustomLSTM, self).__init__()
self.input_dim = input_dim
self.hidden_dim = hidden_dim
self.num_layers = num_layers
self.batch_first = batch_first
# 创建多个 LSTM Cell
self.cells = nn.ModuleList([
CustomLSTMCell(input_dim if l == 0 else hidden_dim, hidden_dim)
for l in range(num_layers)
])
def forward(self, x, hidden=None):
if self.batch_first:
x = x.transpose(0, 1) # 转为 [seq_len, batch_size, input_dim]
seq_len, batch_size, _ = x.size()
# 初始化隐藏状态
if hidden is None:
h_0 = torch.zeros(self.num_layers, batch_size, self.hidden_dim, device=x.device)
c_0 = torch.zeros(self.num_layers, batch_size, self.hidden_dim, device=x.device)
hidden = (h_0, c_0)
h_prev, c_prev = hidden
# 逐层处理
layer_output = x
for layer in range(self.num_layers):
outputs = []
h = h_prev[layer]
c = c_prev[layer]
for t in range(seq_len):
h, c = self.cells[layer](layer_output[t], (h, c))
outputs.append(h.unsqueeze(0))
layer_output = torch.cat(outputs, dim=0)
h_prev[layer] = h # 最后一个时间步的隐藏状态
if self.batch_first:
layer_output = layer_output.transpose(0, 1) # 转回 [batch_size, seq_len, hidden_dim]
return layer_output, (h_prev, c_prev)
class LSTMModel(nn.Module):
def __init__(self, input_dim=10, hidden_dim=64, num_layers=2, num_classes=3):
super(LSTMModel, self).__init__()
# 可以通过nn.LSTM调用PyTorch中的LSTM实现
# self.lstm = nn.LSTM(input_dim, hidden_dim, num_layers, batch_first=True, bidirectional=False)
self.lstm = CustomLSTM(input_dim, hidden_dim, num_layers, batch_first=True)
self.classifier = nn.Sequential(
nn.Linear(hidden_dim, num_classes)
)
def forward(self, x):
# x shape: (batch_size, seq_len, input_dim)
out, _ = self.lstm(x) # out shape: (batch_size, seq_len, hidden_dim)
out = out[:, -1, :] # 取最后一个时刻的输出作为分类依据
out = self.classifier(out)
return out
# ================== Step 3: 工具函数与配置 ==================
def get_dataloaders(batch_size=32):
train_dataset = SequenceDataset()
val_dataset = SequenceDataset() # 实际中应使用真实验证集
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=batch_size)
return train_loader, val_loader
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 train_one_epoch(model, dataloader, criterion, optimizer, device, epoch, writer):
model.train()
running_loss = 0.0
correct = 0
total = 0
for i, (inputs, labels) in enumerate(dataloader):
inputs, labels = inputs.to(device), labels.to(device)
optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
running_loss += loss.item()
_, predicted = outputs.max(1)
total += labels.size(0)
correct += predicted.eq(labels).sum().item()
if i % 10 == 9:
avg_loss = running_loss / 10
acc = 100. * correct / total
logging.info(f"Epoch [{epoch}] Batch [{i}] Loss: {avg_loss:.4f}, Acc: {acc:.2f}%")
writer.add_scalar('Loss/train', avg_loss, epoch * len(dataloader) + i)
running_loss = 0.0
correct = 0
total = 0
@torch.no_grad()
def validate(model, dataloader, criterion, device, epoch, writer):
model.eval()
correct = 0
total = 0
val_loss = 0.0
for inputs, labels in dataloader:
inputs, labels = inputs.to(device), labels.to(device)
outputs = model(inputs)
loss = criterion(outputs, labels)
val_loss += loss.item()
_, predicted = outputs.max(1)
total += labels.size(0)
correct += predicted.eq(labels).sum().item()
acc = 100. * correct / total
avg_loss = val_loss / len(dataloader)
logging.info(f"Epoch [{epoch}] Validation Acc: {acc:.2f}%, Loss: {avg_loss:.4f}")
writer.add_scalar('Loss/val', avg_loss, epoch)
writer.add_scalar('Accuracy/val', acc, epoch)
return acc
# ================== Step 5: 主程序入口 ==================
def main():
setup_logger()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
logging.info(f"Using device: {device}")
batch_size = 32
num_epochs = 10
lr = 0.001
train_loader, val_loader = get_dataloaders(batch_size)
model = LSTMModel(input_dim=10, hidden_dim=64, num_layers=2, num_classes=3).to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=lr)
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.1)
writer = SummaryWriter(log_dir="runs")
best_acc = 0.0
save_dir = "checkpoints"
os.makedirs(save_dir, exist_ok=True)
for epoch in range(num_epochs):
logging.info(f"\nEpoch [{epoch+1}/{num_epochs}]")
train_one_epoch(model, train_loader, criterion, optimizer, device, epoch+1, writer)
val_acc = validate(model, val_loader, criterion, device, epoch+1, writer)
scheduler.step()
# 保存最佳模型
if val_acc > best_acc:
best_acc = val_acc
torch.save(model.state_dict(), os.path.join(save_dir, f"best_model.pth"))
logging.info(f"Saved best model with accuracy: {best_acc:.2f}%")
writer.close()
logging.info("Training complete.")
if __name__ == "__main__":
main()关键要点
| 要点 | 说明 |
|---|---|
| 自定义 LSTM Cell | 手动实现遗忘门、输入门、输出门的计算逻辑 |
| 多层 LSTM | 通过 nn.ModuleList 堆叠多个 LSTM Cell |
| 输出选择 | 取最后一个时间步的输出 out[:, -1, :] 作为分类依据 |
| PyTorch 内置 LSTM | 可直接使用 nn.LSTM(input_dim, hidden_dim, num_layers, batch_first=True) |
扩展建议
| 功能 | 描述 |
|---|---|
| 双向 LSTM | 设置 bidirectional=True 同时捕获前后文信息 |
| Attention | 在 LSTM 输出上添加注意力机制 |
| 真实数据 | 替换为 IMDB 情感分析、时间序列预测等真实数据集 |
| 文本分类 | 配合 nn.Embedding 层用于 NLP 任务 |
PyTorch 所有常用神经网络层完整分类汇总
一、卷积层（Convolution Layers）
| 层名 | 全称 | 作用 | 输入形状 | 输出形状 | 示例 |
|---|---|---|---|---|---|
nn.Conv1d | 一维卷积 | 提取序列数据局部特征（如音频） | (B, C_in, L) | (B, C_out, L') | nn.Conv1d(3, 16, kernel_size=5) |
nn.Conv2d | 二维卷积 | 提取图像局部特征 | (B, C_in, H, W) | (B, C_out, H', W') | nn.Conv2d(3, 16, 3) |
nn.Conv3d | 三维卷积 | 处理视频或医学图像（体积数据） | (B, C_in, D, H, W) | (B, C_out, D', H', W') | nn.Conv3d(3, 16, 3) |
nn.ConvTranspose1d | 一维转置卷积 | 上采样（序列） | (B, C_in, L) | (B, C_out, L') | nn.ConvTranspose1d(16, 3, 4, 2, 1) |
nn.ConvTranspose2d | 二维转置卷积 | 上采样（图像） | (B, C_in, H, W) | (B, C_out, H', W') | nn.ConvTranspose2d(16, 3, 4, 2, 1) |
nn.ConvTranspose3d | 三维转置卷积 | 上采样（体积数据） | (B, C_in, D, H, W) | (B, C_out, D', H', W') | nn.ConvTranspose3d(16, 3, 4, 2, 1) |
nn.Unfold | 将张量展开为滑动窗口块 | 辅助自定义卷积实现 | (B, C, H, W) | (B, C×k×k, L) | nn.Unfold(kernel_size=3, stride=1) |
nn.Fold | 将滑动窗口块重新合并为张量 | 与 Unfold 配合使用 | (B, C×k×k, L) | (B, C, H, W) | nn.Fold(output_size=(H,W), kernel_size=3) |
二、池化层（Pooling Layers）
| 层名 | 全称 | 作用 | 输入形状 | 输出形状 | 示例 |
|---|---|---|---|---|---|
nn.MaxPool1d | 一维最大池化 | 序列下采样 | (B, C, L) | (B, C, L') | nn.MaxPool1d(2) |
nn.MaxPool2d | 二维最大池化 | 图像下采样 | (B, C, H, W) | (B, C, H', W') | nn.MaxPool2d(2) |
nn.MaxPool3d | 三维最大池化 | 体积数据下采样 | (B, C, D, H, W) | (B, C, D', H', W') | nn.MaxPool3d(2) |
nn.AdaptiveMaxPool1d | 自适应最大池化（1D） | 固定输出长度 | (B, C, L) | (B, C, out_L) | nn.AdaptiveMaxPool1d(10) |
nn.AdaptiveMaxPool2d | 自适应最大池化（2D） | 固定输出尺寸 | (B, C, H, W) | (B, C, out_H, out_W) | nn.AdaptiveMaxPool2d((7,7)) |
nn.AdaptiveMaxPool3d | 自适应最大池化（3D） | 固定输出体积 | (B, C, D, H, W) | (B, C, out_D, out_H, out_W) | nn.AdaptiveMaxPool3d((4,4,4)) |
nn.AvgPool1d | 一维平均池化 | 平滑特征 | (B, C, L) | (B, C, L') | nn.AvgPool1d(2) |
nn.AvgPool2d | 二维平均池化 | 平滑特征 | (B, C, H, W) | (B, C, H', W') | nn.AvgPool2d(2) |
nn.AvgPool3d | 三维平均池化 | 平滑特征 | (B, C, D, H, W) | (B, C, D', H', W') | nn.AvgPool3d(2) |
nn.AdaptiveAvgPool1d | 自适应平均池化（1D） | 固定输出长度 | (B, C, L) | (B, C, out_L) | nn.AdaptiveAvgPool1d(10) |
nn.AdaptiveAvgPool2d | 自适应平均池化（2D） | 固定输出尺寸 | (B, C, H, W) | (B, C, out_H, out_W) | nn.AdaptiveAvgPool2d((1,1)) |
nn.AdaptiveAvgPool3d | 自适应平均池化（3D） | 固定输出体积 | (B, C, D, H, W) | (B, C, out_D, out_H, out_W) | nn.AdaptiveAvgPool3d((1,1,1)) |
三、激活函数层（Activation Functions）
| 层名 | 公式 | 作用 | 输入输出相同 |
|---|---|---|---|
nn.ReLU | $f(x)=\max(0,x)$ | 引入非线性，加速收敛 | ✅ |
nn.ReLU6 | $f(x)=\min(\max(0,x),6)$ | 控制输出范围 | ✅ |
nn.LeakyReLU | $f(x)=\begin{cases}x & x>0\ \alpha x & x\leq 0\end{cases}$ | 解决ReLU死亡问题 | ✅ |
nn.PReLU | 可学习参数的LeakyReLU | 自适应控制负值斜率 | ✅ |
nn.ELU | $f(x)=\begin{cases}x & x>0\ \alpha(e^x-1) & x\leq 0\end{cases}$ | 更平滑的负区响应 | ✅ |
nn.Sigmoid | $f(x)=\frac{1}{1+e^{-x}}$ | 输出概率分布（二分类） | ✅ |
nn.Tanh | $f(x)=\tanh(x)$ | 输出 [-1,1] 区间 | ✅ |
nn.Hardtanh | 分段线性近似Tanh | 减少计算开销 | ✅ |
nn.Softplus | $f(x)=\ln(1+e^x)$ | 平滑ReLU替代 | ✅ |
nn.Softsign | $f(x)=\frac{x}{1+\lvert x\rvert}$ | - | ✅ |
nn.Threshold | 低于阈值设为某值 | 简单二值化操作 | ✅ |
nn.Softmax | $f(x_i)=\frac{e^{x_i}}{\sum_j e^{x_j}}$ | 多分类输出归一化 | ❌（dim需指定） |
nn.LogSoftmax | $\log(\text{Softmax}(x))$ | 用于NLLLoss优化 | ❌ |
nn.GELU | 高斯误差线性单元 | Transformer 中常用激活函数 | ✅ |
nn.Mish | $f(x)=x\cdot\tanh(\ln(1+e^x))$ | 新型激活函数 | ✅ |
nn.SiLU / Swish | $f(x)=x\cdot\sigma(x)$ | 自门控激活函数 | ✅ |
四、归一化层（Normalization Layers）
| 层名 | 全称 | 作用 | 输入形状 | 示例 |
|---|---|---|---|---|
nn.BatchNorm1d | 一维批归一化 | 对每个通道做标准化 | (B, C, L) | nn.BatchNorm1d(16) |
nn.BatchNorm2d | 二维批归一化 | 对图像做标准化 | (B, C, H, W) | nn.BatchNorm2d(16) |
nn.BatchNorm3d | 三维批归一化 | 对体积数据标准化 | (B, C, D, H, W) | nn.BatchNorm3d(16) |
nn.InstanceNorm1d | 实例归一化（1D） | 每个样本单独标准化 | (B, C, L) | nn.InstanceNorm1d(16) |
nn.InstanceNorm2d | 实例归一化（2D） | 图像风格迁移中常用 | (B, C, H, W) | nn.InstanceNorm2d(16) |
nn.LayerNorm | 层归一化 | 对每个样本的整个层归一化 | 任意形状 | nn.LayerNorm([C, H, W]) |
nn.GroupNorm | 组归一化 | 分组后进行标准化 | (B, C, H, W) | nn.GroupNorm(4, 16) |
nn.LocalResponseNorm | 局部响应归一化 | AlexNet 中使用 | (B, C, H, W) | nn.LocalResponseNorm(2) |
五、Dropout 层（防止过拟合）
| 层名 | 全称 | 作用 | 输入形状 |
|---|---|---|---|
nn.Dropout | 随机失活 | 训练时随机将部分神经元置零 | 任意 |
nn.Dropout2d | 通道级随机失活 | 对整通道随机置零（图像） | (B, C, H, W) |
nn.Dropout3d | 通道级随机失活（3D） | 适用于体积数据 | (B, C, D, H, W) |
nn.AlphaDropout | 保持均值和方差不变的Dropout | 用于Self-Normalizing Networks | 任意 |
六、循环神经网络层（RNN Layers）
| 层名 | 全称 | 作用 | 输入形状 | 输出形状 |
|---|---|---|---|---|
nn.RNN | 循环神经网络 | 处理序列数据 | (B, T, input_size) | (B, T, hidden_size) |
nn.LSTM | 长短期记忆网络 | 解决长依赖问题 | 同上 | 同上 |
nn.GRU | 门控循环单元 | 简化版LSTM | 同上 | 同上 |
nn.LSTMCell | 单步LSTM | 构建自定义RNN | (B, input_size) | (B, hidden_size) |
nn.GRUCell | 单步GRU | 自定义循环逻辑 | 同上 | 同上 |
七、嵌入与查找层（Embedding Layers）
| 层名 | 全称 | 作用 | 输入形状 | 输出形状 |
|---|---|---|---|---|
nn.Embedding | 嵌入层 | 将离散索引映射为向量 | (B, T) | (B, T, embed_dim) |
nn.EmbeddingBag | 嵌入包层 | 快速聚合词袋模型 | (B, T) 或 (B, T, offsets) | (B, embed_dim) |
八、变换与展平层（Transformation Layers）
| 层名 | 作用 | 输入形状 | 输出形状 |
|---|---|---|---|
nn.Flatten | 将多维张量展平 | (B, C, H, W) | (B, C*H*W) |
nn.Unflatten | 反展平 | (B, N) | (B, C, H, W) |
nn.Upsample / F.interpolate | 图像上采样 | (B, C, H, W) | (B, C, H', W') |
nn.Identity | 恒等映射 | 占位符或条件分支 | 任意 |
九、容器类层（Container Modules）
| 层名 | 作用 | 说明 |
|---|---|---|
nn.Sequential | 顺序执行多个层 | nn.Sequential(Conv, ReLU, Pool) |
nn.ModuleList | 存储子模块列表 | 支持迭代和索引访问 |
nn.ModuleDict | 存储命名子模块 | {'conv': Conv(), 'fc': Linear()} |
nn.Parameter | 存储参数 | nn.Parameter(torch.randn(dim)) 定义一个可学习的权重向量 |
nn.ParameterList | 存储参数列表 | 用于自定义层 |
nn.ParameterDict | 存储命名参数字典 | 同上 |
十、其他功能层
| 层名 | 作用 | 输入输出 |
|---|---|---|
nn.ConstantPad1d/2d/3d | 常数填充 | 改变尺寸 |
nn.ReflectionPad2d | 反射填充（图像处理） | 同上 |
nn.ReplicationPad2d | 边缘复制填充 | 同上 |
nn.ZeroPad2d | 零填充 | 同上 |
nn.Softmax | 归一化为概率分布 | (B, N) → (B, N) |
nn.PairwiseDistance | 计算两个向量之间的距离 | (B, D) × (B, D) → (B,) |
nn.CosineSimilarity | 计算余弦相似度 | 同上 |
十一、Transformer 相关层
| 层名 | 全称 | 作用 | 输入形状 | 输出形状 | 示例 |
|---|---|---|---|---|---|
nn.MultiheadAttention | 多头注意力机制 | 并行计算多个注意力头 | (L, B, E) 或 (B, L, E) | (L, B, E) / (B, L, E) | nn.MultiheadAttention(embed_dim=512, num_heads=8) |
nn.TransformerEncoderLayer | Transformer 编码器层 | 包含自注意力和前馈网络 | (S, B, E) | (S, B, E) | nn.TransformerEncoderLayer(d_model=512, nhead=8) |
nn.TransformerEncoder | Transformer 编码器堆叠 | 堆叠多个编码器层 | (S, B, E) | (S, B, E) | nn.TransformerEncoder(encoder_layer, num_layers=6) |
nn.TransformerDecoderLayer | Transformer 解码器层 | 自注意 + 编码器输出交叉注意 | (T, B, E) | (T, B, E) | nn.TransformerDecoderLayer(d_model=512, nhead=8) |
nn.TransformerDecoder | Transformer 解码器堆叠 | 堆叠多个解码器层 | (T, B, E) | (T, B, E) | nn.TransformerDecoder(decoder_layer, num_layers=6) |
nn.Transformer | 完整的 Transformer 模型 | 包含编码器+解码器 | (S, B, E) & (T, B, E) | (T, B, E) | nn.Transformer(d_model=512, nhead=8) |
说明：
L表示源序列长度，T是目标序列长度，B是 batch size，E是嵌入维度。 默认输入格式是(sequence_length, batch_size, embedding_dim)，如果使用batch_first=True，则变为(batch_size, sequence_length, embedding_dim)。
示例 1：多头注意力机制
attn = nn.MultiheadAttention(embed_dim=512, num_heads=8)
query = key = value = torch.randn(10, 32, 512) # (seq_len, batch_size, embed_dim)
attn_output, _ = attn(query, key, value)
print(attn_output.shape) # 输出: (10, 32, 512)示例 2：构建一个完整的 Transformer 编码器
encoder_layer = nn.TransformerEncoderLayer(d_model=512, nhead=8)
transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=6)
src = torch.randn(20, 32, 512) # (src_len, batch_size, d_model)
output = transformer_encoder(src)
print(output.shape) # 输出: (20, 32, 512)十二、损失函数层（Loss Functions）
| 层名 | 全称 | 作用 | 输入形状 | 输出形状 | 示例 |
|---|---|---|---|---|---|
nn.MSELoss | 均方误差损失 | 回归任务常用 | (N, *), (N, *) | 标量 | loss = nn.MSELoss() |
nn.L1Loss | 平均绝对误差损失 | 回归任务，对异常值更鲁棒 | 同上 | 标量 | loss = nn.L1Loss() |
nn.SmoothL1Loss | 平滑L1损失 | 回归任务中结合MSE和MAE优点 | 同上 | 标量 | loss = nn.SmoothL1Loss() |
nn.CrossEntropyLoss | 交叉熵损失 | 多分类任务（自动包含 Softmax） | (N, C), (N) | 标量 | loss = nn.CrossEntropyLoss() |
nn.NLLLoss | 负对数似然损失 | 多分类任务（需先 Softmax + log） | (N, C), (N) | 标量 | loss = nn.NLLLoss() |
nn.BCELoss | 二元交叉熵损失 | 二分类问题 | (N, *), (N, *) | 标量 | loss = nn.BCELoss() |
nn.BCEWithLogitsLoss | BCE + Sigmoid | 更稳定，推荐用于二分类 | (N, *), (N, *) | 标量 | loss = nn.BCEWithLogitsLoss() |
nn.KLDivLoss | KL 散度损失 | 分布差异衡量 | (N, C), (N, C) | 标量 | loss = nn.KLDivLoss() |
nn.MarginRankingLoss | 排序损失 | 学习排序模型 | (N), (N), (N) | 标量 | loss = nn.MarginRankingLoss(margin=1.0) |
nn.HingeEmbeddingLoss | 支持向量机风格损失 | 度量学习、相似性判断 | (N, D), (N) | 标量 | loss = nn.HingeEmbeddingLoss() |
nn.CosineEmbeddingLoss | 余弦相似性损失 | 判断两个向量是否相似 | (N, D), (N, D), (N) | 标量 | loss = nn.CosineEmbeddingLoss() |
nn.TripletMarginLoss | 三元组损失 | 度量学习、人脸识别 | (N, D), (N, D), (N, D) | 标量 | loss = nn.TripletMarginLoss(margin=1.0) |
nn.PoissonNLLLoss | 泊松负对数似然损失 | 适用于泊松分布建模 | (N, *), (N, *) | 标量 | loss = nn.PoissonNLLLoss() |
nn.CTCLoss | 连接时序分类损失 | 语音识别、OCR 等 | (T, N, C), (N), (T, N) | 标量 | loss = nn.CTCLoss(blank=0) |
nn.GaussianNLLLoss | 高斯负对数似然损失 | 概率回归任务 | (N, *), (N, *), (N, *) | 标量 | loss = nn.GaussianNLLLoss() |
nn.SoftMarginLoss | 多标签二分类损失 | 多个类别可同时为真 | (N, C), (N, C) | 标量 | loss = nn.SoftMarginLoss() |
nn.MultiLabelSoftMarginLoss | 多标签分类损失 | 多标签分类任务 | (N, C), (N, C) | 标量 | loss = nn.MultiLabelSoftMarginLoss() |
nn.MultiMarginLoss | 多类支持向量机损失 | 传统 SVM 扩展 | (N, C), (N) | 标量 | loss = nn.MultiMarginLoss() |
示例 1：交叉熵损失（多分类）
criterion = nn.CrossEntropyLoss()
inputs = torch.randn(3, 5) # (batch_size, num_classes)
targets = torch.empty(3, dtype=torch.long).random_(5) # (batch_size,)
loss = criterion(inputs, targets)
print(loss.item())示例 2：二分类 BCEWithLogitsLoss
bce_criterion = nn.BCEWithLogitsLoss()
logits = torch.randn(4, 1) # (batch_size, 1)
targets = torch.empty(4, 1).random_(2) # (batch_size, 1), 0 or 1
loss = bce_criterion(logits, targets)
print(loss.item())示例 3：CTC Loss（语音识别）
ctc_loss = nn.CTCLoss()
log_probs = torch.randn(50, 16, 20).log_softmax(2) # (T, N, C)
targets = torch.randint(1, 20, (16, 30)) # (N, S)
input_lengths = torch.full((16,), 50, dtype=torch.long)
target_lengths = torch.randint(10, 30, (16,), dtype=torch.long)
loss = ctc_loss(log_probs, targets, input_lengths, target_lengths)
print(loss.item())十三、优化器（Optimizer）
| 优化器名称 | 接口 | 特点 | 适用场景 | 使用示例 |
|---|---|---|---|---|
| SGD（随机梯度下降） | optim.SGD(params, lr=<required>, momentum=0, dampening=0, weight_decay=0, nesterov=False) | 最基础的优化方法；支持动量加速收敛 | 广泛适用于各种问题，尤其是大规模数据集 | optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9) |
| Adam | optim.Adam(params, lr=0.001, betas=(0.9, 0.999), eps=1e-08, weight_decay=0, amsgrad=False) | 自适应学习率，结合了AdaGrad和RMSProp的优点 | 大多数深度学习任务，默认选择之一 | optimizer = optim.Adam(model.parameters(), lr=0.001) |
| RMSprop | optim.RMSprop(params, lr=0.01, alpha=0.99, eps=1e-08, weight_decay=0, momentum=0, centered=False) | 通过平方梯度的移动平均值来归一化学习率 | 循环神经网络(RNN)，非凸优化问题 | optimizer = optim.RMSprop(model.parameters(), lr=0.01, alpha=0.9) |
| Adagrad | optim.Adagrad(params, lr=0.01, lr_decay=0, weight_decay=0, initial_accumulator_value=0, eps=1e-10) | 学习率基于参数更新历史自适应调整 | 稀疏数据，如自然语言处理(NLP)任务 | optimizer = optim.Adagrad(model.parameters(), lr=0.01) |
| Adadelta | optim.Adadelta(params, lr=1.0, rho=0.9, eps=1e-6, weight_decay=0) | Adagrad 的改进版，减少对初始学习率的选择敏感性 | 不需要手动设置学习率的情况 | optimizer = optim.Adadelta(model.parameters()) |
| AdamW | optim.AdamW(params, lr=0.001, betas=(0.9, 0.999), eps=1e-8, weight_decay=0.01, amsgrad=False) | Adam 的修正版本，更好地处理权重衰减 | 防止过拟合，提高泛化能力 | optimizer = optim.AdamW(model.parameters(), lr=0.001, weight_decay=0.01) |
| SparseAdam | optim.SparseAdam(params, lr=0.001, betas=(0.9, 0.999), eps=1e-8) | 专为稀疏张量设计的 Adam 版本 | 只有当梯度是稀疏时使用 | optimizer = optim.SparseAdam(model.parameters(), lr=0.001) |
| ASGD | optim.ASGD(params, lr=0.01, lambd=0.0001, alpha=0.75, t0=1000000.0, weight_decay=0) | 长期运行任务中使用的一种变体，通过平均多个迭代步骤中的参数来改善稳定性 | 需要长期稳定性的任务 | optimizer = optim.ASGD(model.parameters(), lr=0.01) |
| LBFGS | optim.LBFGS(params, lr=1, max_iter=20, max_eval=None, tolerance_grad=1e-7, tolerance_change=1e-9, history_size=100, line_search_fn=None) | 基于拟牛顿法的优化器 | 小规模问题或需要高精度解的任务 | optimizer = optim.LBFGS(model.parameters(), lr=1) |
参数说明：
params：模型参数，通常可以通过model.parameters()获取lr(learning rate)：学习率，决定了步长大小，太大会导致不收敛，太小则训练缓慢weight_decay：权重衰减项，用于正则化，防止过拟合momentum/nesterov：动量技术，帮助加速SGD在相关方向上的进展，并抑制震荡
十四、学习率调度器（Learning Rate Scheduler）
| 名称 | 接口 | 特点 | 适用场景 | 使用示例 |
|---|---|---|---|---|
| StepLR | torch.optim.lr_scheduler.StepLR(optimizer, step_size, gamma=0.1) | 每隔固定 epoch 数衰减学习率（乘以 gamma） | 固定周期调整 LR，简单有效 | scheduler = StepLR(optimizer, step_size=30, gamma=0.1) |
| MultiStepLR | torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones, gamma) | 在指定的多个 epoch 点衰减 LR | 自定义多个阶段的 LR 下降 | scheduler = MultiStepLR(optimizer, milestones=[50, 80], gamma=0.1) |
| ExponentialLR | torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma) | 每个 epoch 都按指数衰减 LR（乘以 gamma） | 连续缓慢降低 LR | scheduler = ExponentialLR(optimizer, gamma=0.95) |
| CosineAnnealingLR | torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max, eta_min=0) | 学习率按余弦函数变化，先下降后上升（一个周期） | 图像分类等任务，模拟温度退火 | scheduler = CosineAnnealingLR(optimizer, T_max=50) |
| ReduceLROnPlateau | torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.1, patience=10) | 根据监控指标（如验证 loss）自动调整 LR | 验证 loss 不再改善时降低 LR | scheduler = ReduceLROnPlateau(optimizer, 'min', patience=5) |
| CyclicLR | torch.optim.lr_scheduler.CyclicLR(optimizer, base_lr, max_lr, step_size_up=2000) | 在两个学习率之间循环变化 | 快速训练、探索 LR 敏感性 | scheduler = CyclicLR(optimizer, base_lr=0.001, max_lr=0.01, step_size_up=100) |
| OneCycleLR | torch.optim.lr_scheduler.OneCycleLR(optimizer, max_lr, steps_per_epoch, epochs) | 一个周期内先升后降的学习率策略 | 大多数现代 CNN 训练推荐 | scheduler = OneCycleLR(optimizer, max_lr=0.01, steps_per_epoch=len(train_loader), epochs=10) |
| LambdaLR | torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda) | 自定义学习率调整函数 | 完全灵活控制 LR 调整方式 | lambda epoch: 0.95 ** epoch |
| MultiplicativeLR | torch.optim.lr_scheduler.MultiplicativeLR(optimizer, lr_lambda) | 每次调用 step() 时乘以一个因子 | 简单自定义 LR 缩放 | lambda epoch: 0.9 |
调度器选择指南
| 目标 | 推荐调度器 |
|---|---|
| 简单实用、逐步下降 | StepLR, MultiStepLR |
| 更精细地控制 LR 变化 | LambdaLR, MultiplicativeLR |
| 验证集指标驱动 | ReduceLROnPlateau |
| 最新推荐、性能最好 | OneCycleLR |
| 循环学习率、尝试不同范围 | CyclicLR |
| 模拟退火效果 | CosineAnnealingLR |
使用示例
import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim.lr_scheduler import StepLR
# 模型和优化器
model = nn.Linear(10, 2)
optimizer = optim.Adam(model.parameters(), lr=0.1)
# 设置调度器
scheduler = StepLR(optimizer, step_size=3, gamma=0.1)
for epoch in range(10):
for inputs, labels in dataloader:
optimizer.zero_grad()
outputs = model(inputs)
loss = loss_fn(outputs, labels)
loss.backward()
optimizer.step()
scheduler.step() # 更新学习率总结
| 类别 | 数量 |
|---|---|
| 卷积层（Convolution Layers） | ~8 |
| 池化层（Pooling Layers） | ~12 |
| 激活函数层（Activation Functions） | ~16 |
| 归一化层（Normalization Layers） | ~8 |
| Dropout 层（防止过拟合） | ~4 |
| 循环神经网络层（RNN Layers） | ~5 |
| 嵌入与查找层（Embedding Layers） | ~2 |
| 变换与展平层（Transformation Layers） | ~4 |
| 容器类层（Container Modules） | ~5 |
| 其他功能层 | ~7 |
| Transformer 层 | ~7 |
| 损失函数层（Loss Layers） | ~19 |
PyTorch 自定义层
一、为什么需要自定义层？
虽然 PyTorch 提供了大量内置层（如 nn.Conv2d, nn.Linear, nn.LSTM 等），但在以下情况下你可能需要自定义层：
- 实现论文中提出的新模块
- 构建非标准网络结构
- 封装特定的前向逻辑以提高代码可读性
- 添加可学习参数（权重 + 偏置）
- 控制梯度传播行为
二、如何在 PyTorch 中自定义一个层？
- 步骤 1：继承
torch.nn.Module，这是所有 PyTorch 层的基类。 - 步骤 2：定义
__init__方法，用于初始化可学习参数和子模块（如 Conv、Linear 等）。 - 步骤 3：实现
forward方法，定义该层的前向传播逻辑。
示例 1：最简单的自定义层 —— 自定义线性变换层
import torch
import torch.nn as nn
class MyLinear(nn.Module):
def __init__(self, input_dim, output_dim):
super(MyLinear, self).__init__()
# 定义可学习参数
self.weight = nn.Parameter(torch.randn(output_dim, input_dim))
self.bias = nn.Parameter(torch.randn(output_dim))
def forward(self, x):
# 实现矩阵乘法 y = Wx + b
return x @ self.weight.t() + self.bias使用示例：
linear = MyLinear(10, 5)
x = torch.randn(3, 10) # batch_size=3, input_dim=10
y = linear(x)
print(y.shape) # 输出: torch.Size([3, 5])示例 2：带激活函数的复合层
class MyDenseBlock(nn.Module):
def __init__(self, in_features, out_features, activation=nn.ReLU()):
super(MyDenseBlock, self).__init__()
self.linear = nn.Linear(in_features, out_features)
self.activation = activation
def forward(self, x):
return self.activation(self.linear(x))使用示例：
block = MyDenseBlock(20, 10, activation=nn.ReLU())
x = torch.randn(4, 20)
y = block(x)
print(y.shape) # 输出: torch.Size([4, 10])示例 3：使用已有层构建更复杂的自定义块（如 ResNet Block）
class ResidualBlock(nn.Module):
def __init__(self, channels):
super(ResidualBlock, self).__init__()
self.conv1 = nn.Conv2d(channels, channels, kernel_size=3, padding=1)
self.conv2 = nn.Conv2d(channels, channels, kernel_size=3, padding=1)
self.relu = nn.ReLU()
def forward(self, x):
residual = x
x = self.relu(self.conv1(x))
x = self.conv2(x)
x += residual # 残差连接
x = self.relu(x)
return x三、自定义层的关键要素总结
| 元素 | 说明 |
|---|---|
__init__ | 初始化参数、子模块 |
forward() | 定义前向传播逻辑 |
nn.Parameter | 标记为可学习参数，自动加入优化器 |
子模块（如 nn.Conv2d） | 可直接调用 PyTorch 内置层 |
| 支持 GPU 加速 | 调用 .to(device) 即可 |
四、进阶技巧
1. 添加不可学习的缓冲区（buffer）
比如你需要保存一些不需要梯度更新的张量（如位置编码）：
def __init__(self):
super().__init__()
self.register_buffer('position_ids', torch.arange(512).expand((1, -1)))2. 支持任意输入维度（通用层）
如果你希望你的层适用于不同维度的数据（如 1D、2D、3D），可以使用 F.linear, F.convNd 等函数代替具体层。
3. 自定义反向传播（高级）
如果你需要完全控制梯度计算，可以使用 torch.autograd.Function 来自定义前向和反向传播。
class MyFunction(torch.autograd.Function):
@staticmethod
def forward(ctx, input):
ctx.save_for_backward(input)
return input.clamp(min=0)
@staticmethod
def backward(ctx, grad_output):
input, = ctx.saved_tensors
grad_input = grad_output.clone()
grad_input[input < 0] = 0
return grad_input然后可以在你的 nn.Module 中调用它：
class MyReLU(nn.Module):
def forward(self, x):
return MyFunction.apply(x)五、验证你的自定义层是否正常工作
建议在实现后做以下测试：
- 输入输出形状是否正确
- 是否支持 GPU (
layer.to("cuda")) - 是否能正常进行反向传播（不报错）
- 参数是否被优化器识别并更新
layer = MyLinear(10, 5)
optimizer = torch.optim.Adam(layer.parameters(), lr=1e-3)
loss = layer(torch.randn(2, 10)).sum()
loss.backward()
optimizer.step()六、封装成 Sequential 或嵌入模型中
你可以像使用任何其他 PyTorch 层一样将自定义层放入 nn.Sequential 或组合进模型中：
model = nn.Sequential(
MyLinear(10, 64),
nn.ReLU(),
MyLinear(64, 10)
)总结：PyTorch 自定义层的”黄金模板”
import torch
import torch.nn as nn
class YourCustomLayer(nn.Module):
def __init__(self, *args, **kwargs):
super(YourCustomLayer, self).__init__()
# 初始化参数或子模块
def forward(self, x):
# 前向传播逻辑
return x数据集与转换工具
PyTorch 图像数据集 — torchvision
| 类型 | 接口 | 作用 | 示例 |
|---|---|---|---|
torchvision.datasets | MNIST | 加载手写数字数据集（灰度图） | datasets.MNIST(root='./data', train=True, download=True) |
torchvision.datasets | CIFAR10 | 加载CIFAR-10数据集（彩色图） | datasets.CIFAR10(root='./data', train=True, download=True) |
torchvision.datasets | ImageFolder | 从文件夹中加载图像数据集 | datasets.ImageFolder(root='./flowers', transform=transform) |
torchvision.datasets | VOCDetection | 加载PASCAL VOC 数据集中的检测任务部分 | datasets.VOCDetection(root='./data', year='2012', image_set='train', download=True) |
torchvision.datasets | VOCSegmentation | 加载PASCAL VOC 数据集中的分割任务部分 | datasets.VOCSegmentation(root='./data', year='2012', image_set='train', download=True) |
torchvision.models | resnet18 | 预训练的ResNet-18模型 | models.resnet18(pretrained=True) |
torchvision.models | vgg16 | 预训练的VGG-16模型 | models.vgg16(pretrained=True) |
torchvision.models | alexnet | 预训练的AlexNet模型 | models.alexnet(pretrained=True) |
torchvision.models | inception_v3 | 预训练的Inception v3模型 | models.inception_v3(pretrained=True) |
torchvision.models | mobilenet_v2 | 轻量级的MobileNet V2模型 | models.mobilenet_v2(pretrained=True) |
torchvision.transforms | ToTensor | 将PIL图像或NumPy数组转换为PyTorch张量，并将像素值从[0, 255]缩放到[0.0, 1.0] | transforms.ToTensor() |
torchvision.transforms | Normalize | 使用给定的均值和标准差对每个通道进行标准化 | transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) |
torchvision.transforms | Resize | 调整输入图像的大小到指定尺寸 | transforms.Resize((256, 256)) |
torchvision.transforms | CenterCrop | 在中心位置裁剪出指定大小的图像 | transforms.CenterCrop(224) |
torchvision.transforms | RandomCrop | 随机位置裁剪出指定大小的图像 | transforms.RandomCrop(224, padding=4) |
torchvision.transforms | RandomHorizontalFlip | 以给定的概率随机水平翻转图像 | transforms.RandomHorizontalFlip(p=0.5) |
torchvision.transforms | RandomVerticalFlip | 以给定的概率随机垂直翻转图像 | transforms.RandomVerticalFlip(p=0.5) |
torchvision.transforms | ColorJitter | 随机改变图片的亮度、对比度、饱和度和色调 | transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.2) |
torchvision.transforms | RandomRotation | 随机旋转图像一定角度 | transforms.RandomRotation(30) |
torchvision.transforms | Grayscale | 将图像转换为灰度图 | transforms.Grayscale(num_output_channels=1) |
torchvision.transforms | RandomGrayscale | 随机概率地将图像转换为灰度图 | transforms.RandomGrayscale(p=0.1) |
torchvision.transforms | Pad | 对图像四周填充指定数量的像素 | transforms.Pad(padding=4, fill=0, padding_mode='constant') |
torchvision.transforms | RandomAffine | 随机仿射变换（包括旋转、平移、缩放等） | transforms.RandomAffine(degrees=30, translate=None, scale=None, shear=None) |
torchvision.transforms | RandomPerspective | 随机透视变换 | transforms.RandomPerspective(distortion_scale=0.5, p=0.5, interpolation=Image.BILINEAR) |
torchvision.transforms | RandomResizedCrop | 先随机裁剪再调整大小，常用于训练阶段增加数据多样性 | transforms.RandomResizedCrop(size=224, scale=(0.08, 1.0), ratio=(0.75, 1.3333333333333333)) |
torchvision.transforms | FiveCrop | 在四个角和中心位置各裁剪一个指定大小的子图 | transforms.FiveCrop(size=224) |
torchvision.transforms | TenCrop | 结合FiveCrop加上其水平翻转版本 | transforms.TenCrop(size=224) |
PyTorch 文本数据集 — torchtext
| 类型 | 接口 | 用途 | 示例 |
|---|---|---|---|
torchtext.datasets | AG_NEWS | 加载 AG 新闻数据集 | from torchtext.datasets import AG_NEWS; train_iter, test_iter = AG_NEWS() |
torchtext.datasets | AmazonReviewFull | 加载亚马逊评论全量数据集 | AmazonReviewFull(root='./data') |
torchtext.datasets | DBpedia | DBpedia 分类数据集 | DBpedia(root='./data') |
torchtext.datasets | IMDB | 加载电影评论情感分析数据集 | IMDB(root='./data') |
torchtext.datasets | SST2 | Stanford Sentiment Treebank 数据集 | SST2(root='./data') |
torchtext.datasets | WikiText2 | 小型维基百科语言模型数据集 | WikiText2(root='./data') |
torchtext.datasets | PennTreebank | 标准语言模型训练数据集 | PennTreebank(root='./data') |
torchtext.datasets | Multi30k | 多语言翻译数据集 | Multi30k(root='./data', language_pair=('de', 'en')) |
torchtext.transforms | ToTensor | 将序列转换为张量 | transforms.ToTensor(padding_value=1) |
torchtext.transforms | Truncate | 截断文本长度 | transforms.Truncate(max_length=512) |
torchtext.transforms | AddToken | 在开头或结尾添加特殊 token | AddToken(token=2, begin=True) |
torchtext.transforms | LabelToIndex | 将标签字符串映射为索引 | LabelToIndex(label_names=['pos', 'neg']) |
torchtext.transforms | Sequential | 组合多个 transform | Sequential(tokenizer, vocab_func, to_tensor) |
torchtext.transforms | SentencePieceTokenizer | 使用 SentencePiece 模型对文本进行分词 | tokenizer = SentencePieceTokenizer(model_path="path/to/spm.model") |
torchtext.transforms | BERTTokenizer | 使用预训练的 BERT 模型进行分词 | tokenizer = BERTTokenizer(vocab_path="path/to/vocab.txt", do_lower_case=True) |
torchtext.transforms | VocabTransform | 将词汇表中的词转换为其对应的索引 | vocab_transform = VocabTransform(vocab) |
torchtext.transforms | CharBPETokenizer | 使用字符级别的 Byte-Pair Encoding 方法进行分词 | tokenizer = CharBPETokenizer(merges_file="path/to/merges.txt", vocab_file="path/to/vocab.json") |
torchtext.transforms | GPT2Tokenizer | 使用预训练的 GPT-2 模型进行分词 | tokenizer = GPT2Tokenizer(vocab_file="path/to/gpt2-vocab.json", merges_file="path/to/gpt2-merges.txt") |
torchtext.transforms | WhitespaceTokenizer | 根据空格对文本进行简单分词 | tokenizer = WhitespaceTokenizer() |
torchtext.transforms | RegexTokenizer | 使用正则表达式模式对文本进行分词 | tokenizer = RegexTokenizer(pattern=r'\w+') |
torchtext.transforms | PretrainedSPTokenizer | 加载预训练的 SentencePiece 模型用于分词 | tokenizer = PretrainedSPTokenizer("path/to/pretrained/model") |
torchtext.vocab | Vocab | 创建自定义词汇表 | from torchtext.vocab import Vocab; vocab = Vocab(counter) |
torchtext.vocab | GloVe | 加载预训练 GloVe 向量 | GloVe(name='6B', dim=100) |
torchtext.vocab | FastText | 加载 FastText 预训练向量 | FastText(language='en') |
torchtext.vocab | build_vocab_from_iterator | 从迭代器构建词汇表 | build_vocab_from_iterator(yield_tokens()) |
torchtext.data.utils | get_tokenizer | 获取分词器 | get_tokenizer("basic_english") |
torchtext.data.utils | yield_tokens | 从数据集中生成 token 序列 | for tokens in yield_tokens(train_dataset): ... |
torchtext.functional | ngrams_func | 构建 n-gram 特征 | ngrams_func([1, 2, 3], ngrams=2) |
torchtext.functional | numericalize_tokens_from_vocabulary | 将 token 转为索引 | numericalize_tokens_from_vocabulary(vocab, tokenizer(text)) |
torchtext.dataloader | DataLoader | 支持 collate_fn 的数据加载器 | DataLoader(dataset, batch_size=8, collate_fn=collate_batch) |
torchtext.models | RobertaClassificationHead | RoBERTa 分类头 | RobertaClassificationHead(num_classes=2) |
torchtext.models | TextTransformerMasking | Transformer 掩码模型 | TextTransformerMasking(...) |
torchtext.models | IntraAttention | 带注意力机制的句子分类模型 | IntraAttention(...) |
torchtext.models.classification | text_classification_model | 快速构建分类模型 | text_classification_model(name="AG_NEWS", num_classes=4) |
torchtext.models.roberta | RobertaModel | RoBERTa 主干网络 | RobertaModel.from_pretrained('roberta.base') |
torchtext.models.roberta | RobertaEncoder | RoBERTa 编码器部分 | RobertaEncoder(...) |
torchtext.models.seq2seq | Seq2SeqTransformer | 序列到序列 Transformer 模型 | Seq2SeqTransformer(src_vocab_size, tgt_vocab_size) |
torchtext.data | Dataset | 基础数据集抽象类 | class MyDataset(Dataset): ... |
torchtext.data | IterDataPipe | 可迭代数据管道 | dp = IterableWrapper(data).map(fn).shuffle().batch(32) |
torchtext.data | MapDataPipe | 映射型数据管道 | dp = SequenceWrapper(data).map(fn) |
torchtext.data.functional | to_map_style_dataset | 将可迭代数据集转为 map-style | to_map_style_dataset(iter_dataset) |
torchtext.experimental | torchtext.experimental.vectors | 实验性词向量支持 | from torchtext.experimental.vectors import pretrained_aliases |
torchtext.experimental | torchtext.experimental.models | 实验性模型结构 | from torchtext.experimental.models import bert_base |
PyTorch 音频数据集 — torchaudio
| 类型 | 名称 | 作用 | 示例 |
|---|---|---|---|
torchaudio.datasets | YESICS | 加载 YESICS 数据集 | datasets.YESICS(root='./data', url='MOS-100k') |
torchaudio.datasets | LIBRISPEECH | 加载 LibriSpeech ASR 数据集 | datasets.LIBRISPEECH(root='./data', url='train-clean-100', download=True) |
torchaudio.datasets | SPEECHCOMMANDS | 加载 SpeechCommands v2 数据集 | datasets.SPEECHCOMMANDS(root='./data', url='speech_commands_v0.02', download=True) |
torchaudio.transforms | MelSpectrogram | 计算梅尔频谱图 | MelSpectrogram(sample_rate=16000, n_fft=400, hop_length=160, n_mels=128) |
torchaudio.transforms | AmplitudeToDB | 将幅度转换为分贝 | AmplitudeToDB(stype='power', top_db=80) |
torchaudio.transforms | Resample | 改变音频采样率 | Resample(orig_freq=44100, new_freq=16000) |
torchaudio.transforms | TimeStretch | 音频时间拉伸（保持音调） | TimeStretch(hop_length=512, n_freq=1025) |
torchaudio.transforms | Fade | 淡入淡出效果 | Fade(fade_in_len=1000, fade_out_len=1000, fade_shape='linear') |
torchaudio.transforms | Vad | 声音活动检测 | Vad(sample_rate=16000) |
torchaudio | load | 从文件加载音频 | waveform, sample_rate = torchaudio.load('path/to/audio.wav') |
torchaudio | save | 保存音频到文件 | torchaudio.save('path/to/output.wav', waveform, sample_rate) |
torchaudio | info | 获取音频文件信息 | info = torchaudio.info('path/to/audio.wav') |
torchaudio.functional | compute_deltas | 计算梅尔频谱图的 delta 系数 | compute_deltas(specgram, win_length=5) |
torchaudio.functional | contrast | 应用对比度增强 | contrast(waveform, enhancement_amount=75) |
torchaudio.functional | sliding_window_cmn | 滑动窗口归一化 | sliding_window_cmn(input, cmn_window=600, min_cmn_window=100, center=False, norm_vars=False) |
torchaudio.sox_effects | apply_effects_tensor | 对音频张量应用 SoX 效果链 | effects = [['lowpass', '-1', '300'], ['rate', '16000']]; output, sample_rate = apply_effects_tensor(waveform, sample_rate, effects) |
torchaudio.sox_effects | apply_effects_file | 对音频文件应用 SoX 效果链 | output, sample_rate = apply_effects_file('path/to/input.wav', effects) |
张量操作与连续存储
| 操作 | 作用描述 | 改变物理内存 | 改变读取逻辑 | 是否需要前置 .contiguous() |
|---|---|---|---|---|
.clone() | 创建一个张量的完全副本，包含数据和元信息 | 是 | 否 | 不需要 |
.contiguous() | 返回一个内存布局连续的新张量 | 是 | 否 | 不适用（自身就是用来生成连续张量的） |
.transpose(dim0, dim1) | 交换两个维度的位置 | 否 | 是 | 如果后续要使用 view() 或 reshape()，则需要 |
.permute(*dims) | 对多个维度进行重新排列 | 否 | 是 | 如果后续要使用 view() 或 reshape()，则需要 |
.narrow(dim, start, length) | 在指定维度上创建一个子集视图 | 否 | 是 | 不需要 |
.select(dim, index) | 在指定维度选择特定索引的切片 | 否 | 是 | 不需要 |
.unsqueeze(dim) | 在指定维度插入一个新的维度 | 否 | 是 | 不需要 |
.squeeze(dim) | 移除所有大小为1的维度（除非指定了保留的维度） | 否 | 是 | 不需要 |
.flip(dims) | 沿着指定维度翻转张量 | 否 | 是 | 不需要 |
.rot90(k, dims) | 沿指定轴旋转张量 90 度 | 否 | 是 | 不需要 |
.view(size) | 改变张量的形状 | 是 | 否 | 非连续张量需要前置 .contiguous() |
.reshape(size) | 类似于 .view()，但可以处理非连续张量 | 可能 | 否 | 必要时内部自动调用 .contiguous() |
.to(device) 或 .cuda() | 将张量移动到不同的设备（如从 CPU 到 GPU） | 是 | 否 | 不需要 |
.add(), mul(), div() 等数学运算 | 执行元素级运算 | 是 | 否 | 不需要 |
解释：
- 改变物理内存：表示该操作是否会实际复制数据到新的内存位置。
- 改变读取逻辑：表示该操作是否只是改变了访问张量数据的方式，而不改变其在内存中的实际布局。
- 是否需要前置
.contiguous()：对于某些操作（如.view()），如果当前张量不是连续的（即.is_contiguous()返回False），则需要先调用.contiguous()来确保数据是连续存储的，才能正确执行这些操作。
模型保存与加载
| 保存方法 | 文件格式 | 是否包含模型结构 | 是否推荐使用 | 适用场景 | 保存代码示例 | 加载代码示例 |
|---|---|---|---|---|---|---|
torch.save(model, 'model.pth') | .pth, .pt | ✅ 是（整个模型） | ❌ 否 | 快速保存/加载用于调试或研究环境 | torch.save(model, 'model.pth') | model = torch.load('model.pth') |
torch.save(model.state_dict(), 'model_state.pth') | .pth, .pt | ❌ 否（仅参数） | ✅ 是 | 推荐用于部署、迁移学习、继续训练等场景 | torch.save(model.state_dict(), 'model_state.pth') | model = MyModel(); model.load_state_dict(torch.load('model_state.pth')) |
torch.jit.save(script_model, 'script_model.pt') | .pt, .pth | ✅ 是（序列化模型） | ✅ 是（用于生产部署） | 部署到非 Python 环境（如 C++、移动端） | script_model = torch.jit.script(model); torch.jit.save(script_model, 'script_model.pt') | script_model = torch.jit.load('script_model.pt') |
torch.onnx.export() | .onnx | ✅ 是（标准化模型结构和参数） | ✅ 是（用于跨框架部署） | 转换为 ONNX 格式供 TensorFlow、ONNX Runtime 等使用 | dummy_input = torch.randn(1, 3, 224, 224); torch.onnx.export(model, dummy_input, "model.onnx") | import onnxruntime as ort; session = ort.InferenceSession("model.onnx") |
补充说明
1. 保存整个模型（不推荐）
- 原理：通过 pickle 序列化整个模型对象。
- 缺点：
- 依赖当前的类定义和目录结构。
- 不同版本 PyTorch 可能存在兼容性问题。
- 不适合长期存储或部署。
2. 保存 state_dict（推荐）
- 推荐理由：
- 更加灵活。
- 易于版本控制。
- 支持模型微调、迁移学习等。
- 注意：加载时必须先定义相同的模型结构。
3. TorchScript / ScriptModule（用于部署）
- 适用于需要脱离 Python 环境运行的场景（如 C++、Android、iOS）。
- 可以将模型转换为静态图形式，便于优化和部署。
4. ONNX 导出（跨平台互操作）
- ONNX 是一个开放的神经网络交换格式。
- 可用于在不同框架之间转换模型（如 PyTorch → ONNX → TensorFlow）。
- 支持多种推理引擎（如 ONNX Runtime、TensorRT）。
为什么 TensorFlow 和 PyTorch 的”层”的数量看起来差异大？
总结一句话
PyTorch 更倾向于将各种操作封装为”nn.Module 子类”，而 TensorFlow/Keras 则倾向于通过”函数式 API + 预定义层”来实现功能，导致”显式列出的层”数量不同。
详细解释
1. 设计哲学不同
| 特性 | PyTorch | TensorFlow/Keras |
|---|---|---|
| 面向对象 vs 函数式 | 强调面向对象编程（OO），每个操作都封装成 nn.Module 子类 | 更偏向函数式接口，很多操作是函数而非”层” |
| 动态图 vs 静态图 | 动态计算图（eager execution），便于调试和扩展 | 支持静态图（Graph Execution），也支持 eager 模式 |
| 灵活性 vs 易用性 | 更灵活、适合研究和实验 | 更注重工程化部署、易用性和模型导出 |
因此，在 PyTorch 中，很多操作被封装为独立的类（如 nn.AdaptiveAvgPool2d, nn.InstanceNorm2d, nn.GELU 等），而在 TensorFlow 中，这些可能只是函数（如 tf.nn.relu, tf.nn.gelu）或通过参数控制的通用层。
2. 激活函数等基础操作的封装方式不同
| 类别 | PyTorch | TensorFlow |
|---|---|---|
| 激活函数是否为”层” | 是，例如 nn.ReLU() 是一个类 | 否，通常使用函数如 tf.nn.relu(x) 或 layers.Activation('relu') |
| 是否统一为”Layer”子类 | 是，几乎所有操作都是 nn.Module 的子类 | 否，很多是函数或配置参数 |
这使得在列举”层”时，PyTorch 的数量显得更多，因为每一个小操作都有自己的类。
3. 模块划分粒度不同
| 示例 | PyTorch | TensorFlow |
|---|---|---|
| Dropout 层 | 多种类型：nn.Dropout, nn.Dropout2d, nn.Dropout3d | 只有一个通用 layers.Dropout，自动适应输入维度 |
| BatchNorm 层 | nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d | 一个 layers.BatchNormalization，根据输入自动推断维度 |
| 池化层 | nn.MaxPool1d, nn.MaxPool2d, nn.MaxPool3d | layers.MaxPooling1D/2D/3D，但结构上仍统一处理 |
在 PyTorch 中，为了保持严格的维度一致性，对每种数据结构（1D、2D、3D）都提供了专门的层；而 TensorFlow 使用统一接口自动判断维度。
4. Transformer 相关组件的实现方式
| 方面 | PyTorch | TensorFlow |
|---|---|---|
| MultiheadAttention | 官方提供完整类 nn.MultiheadAttention | 提供 layers.MultiHeadAttention，但不如 PyTorch 完善，部分功能需自定义 |
| Transformer 编码器/解码器 | 作为标准层提供 | 通常需要手动组合多头注意力 + FFN 实现 |
| 位置编码 | 无内置层，需自定义 | 同样需自定义 |
PyTorch 的 nn 模块中包含更多现代架构相关的原生支持，而 TensorFlow 很多功能需要借助 TF Model Garden 或自定义实现。
5. 一些操作在 TensorFlow 中不是”层”，而是函数
| 操作 | PyTorch 层 | TensorFlow 函数 |
|---|---|---|
| Flatten | nn.Flatten() | 一般用 tf.keras.layers.Flatten()，也可直接用 tf.reshape |
| ReLU | nn.ReLU() | tf.nn.relu(x) 或 layers.Activation('relu') |
| Upsample | nn.Upsample() | layers.UpSampling2D() 或 tf.image.resize() |
这意味着 TensorFlow 的”层”列表不会像 PyTorch 那样详尽地列出每一个小功能。
所以，到底哪个框架”功能更多”？
答案是：两者功能几乎一样丰富，只是组织方式不同。
| 维度 | PyTorch | TensorFlow |
|---|---|---|
| 显式”层”的数量 | 更多 | 较少 |
| 操作的灵活性 | 更高（可继承 Layer 自定义） | 更简洁（函数式风格） |
| 研究友好性 | 更适合研究 | 更适合生产部署 |
| 模型导出与部署能力 | 需 ONNX 或 TorchScript | 原生支持 SavedModel、TFLite、TF.js 等 |
举个例子对比
PyTorch 中：每个操作都是一个类实例。
import torch.nn as nn
model = nn.Sequential(
nn.Conv2d(3, 16, 3),
nn.ReLU(),
nn.MaxPool2d(2),
nn.BatchNorm2d(16),
nn.Flatten(),
nn.Linear(16 * 14 * 14, 10)
)TensorFlow/Keras 中：有些操作是函数式封装，比如 Activation('relu')，而不是单独的类。
from tensorflow.keras import layers, models
model = models.Sequential([
layers.Conv2D(16, (3, 3), input_shape=(32, 32, 3)),
layers.Activation('relu'),
layers.MaxPooling2D((2, 2)),
layers.BatchNormalization(),
layers.Flatten(),
layers.Dense(10)
])如何选择？
| 场景 | 推荐框架 |
|---|---|
| 科研、论文复现、快速实验 | PyTorch |
| 工程部署、模型上线、跨平台导出 | TensorFlow |
| 图像识别、CV任务 | 两者均可 |
| NLP、Transformer 架构 | PyTorch 更全面（目前） |
| 分布式训练、TPU 加速 | TensorFlow 更成熟 |
总结：为什么感觉 PyTorch 的层更多？
| 原因 | 说明 |
|---|---|
| 面向对象的设计 | 所有操作都是 nn.Module 的子类，易于统一管理和扩展 |
| 严格区分维度 | 不同维度的操作（1D、2D、3D）都提供独立类 |
| 内置现代架构支持 | Transformer、GELU、Mish 等新型层都有官方实现 |
| 社区贡献活跃 | PyTorch Hub、TorchVision、TorchText 提供大量预定义层 |