### CIFAR-10 断点续训 Python 代码及训练时间统计 以下是基于 TensorFlow 和 PyTorch 的两种实现方式，均支持断点续训功能，并能够统计单个周期的平均训练时间。 --- #### **TensorFlow 实现** 以下是一个完整的 TensorFlow 实现方案，其中包含了断点续训逻辑以及每轮训练的时间记录： ```python import tensorflow as tf import numpy as np import time import os # 超参数设置 train_epochs = 25 batch_size = 50 total_batch = int(len(Xtrain) / batch_size) # 初始化变量用于存储精度和损失 epoch_list = [] accuracy_list = [] loss_list = [] # 定义全局 epoch 变量 epoch = tf.Variable(0, name='epoch', trainable=False) # 创建会话并初始化变量 startTime = time.time() sess = tf.compat.v1.Session() init = tf.compat.v1.global_variables_initializer() sess.run(init) # 设置检查点存储目录 ckpt_dir = "CIFAR_log/" if not os.path.exists(ckpt_dir): os.makedirs(ckpt_dir) # 加载已有检查点 saver = tf.train.Saver(max_to_keep=3) checkpoint = tf.train.latest_checkpoint(ckpt_dir) if checkpoint: saver.restore(sess, checkpoint) start_epoch = sess.run(epoch) else: start_epoch = 0 # 迭代训练过程 for ep in range(start_epoch, train_epochs): avg_loss = 0. avg_acc = 0. for i in range(total_batch): batch_x, batch_y = get_train_batch(i, batch_size) _, l, acc = sess.run([optimizer, loss_function, accuracy], feed_dict={x: batch_x, y: batch_y}) avg_loss += l / total_batch avg_acc += acc / total_batch if i % 100 == 0: print(f'Step {i} finished') # 更新日志 epoch_list.append(ep + 1) loss_list.append(avg_loss) accuracy_list.append(avg_acc) print(f'Train Epoch: {"%02d" % (ep + 1)}, Loss= {"%.6f" % avg_loss}, Accuracy= {"%.4f" % avg_acc}') # 保存检查点 saver.save(sess, ckpt_dir + 'CIFAR10_cnn_model.ckpt', global_step=ep + 1) # 增加当前 epoch 数值 sess.run(epoch.assign(ep + 1)) endTime = time.time() print(f'Average Time per Epoch: {(endTime - startTime) / (ep + 1)} seconds') ``` 此代码实现了断点续训的功能，并通过 `tf.train.Saver` 来管理模型权重的保存与加载[^1]。每次完成一轮训练后都会计算该轮次的平均损失和准确率，并将其存入列表以便后续分析。 --- #### **PyTorch 实现** 对于 PyTorch 用户而言，可以采用如下方法来构建具有断点续训能力的 CIFAR-10 图像分类器： ```python import torch from torch import nn, optim from torchvision import datasets, transforms from torch.utils.data import DataLoader import time import os # 数据预处理 transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,)) ]) train_dataset = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform) test_dataset = datasets.CIFAR10(root='./data', train=False, download=True, transform=transform) train_loader = DataLoader(train_dataset, batch_size=50, shuffle=True) test_loader = DataLoader(test_dataset, batch_size=50, shuffle=False) # 定义网络结构（简单 CNN） class SimpleCNN(nn.Module): def __init__(self): super(SimpleCNN, self).__init__() self.conv_layers = nn.Sequential( nn.Conv2d(in_channels=3, out_channels=32, kernel_size=3), nn.ReLU(), nn.MaxPool2d(kernel_size=2), nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3), nn.ReLU(), nn.MaxPool2d(kernel_size=2) ) self.fc_layers = nn.Sequential( nn.Linear(64 * 6 * 6, 512), nn.ReLU(), nn.Dropout(p=0.5), nn.Linear(512, 10) ) def forward(self, x): x = self.conv_layers(x) x = x.view(-1, 64 * 6 * 6) x = self.fc_layers(x) return x net = SimpleCNN() # 损失函数与优化器 criterion = nn.CrossEntropyLoss() optimizer = optim.Adam(net.parameters(), lr=0.001) # 设备配置 device = torch.device("cuda" if torch.cuda.is_available() else "cpu") net.to(device) # 断点续训路径 model_path = './models/cifar10.pth' start_epoch = 0 if os.path.exists(model_path): checkpoint = torch.load(model_path) net.load_state_dict(checkpoint['model_state_dict']) optimizer.load_state_dict(checkpoint['optimizer_state_dict']) start_epoch = checkpoint['epoch'] + 1 # 训练循环 epochs = 25 training_time_per_epoch = [] for epoch in range(start_epoch, epochs): net.train() running_loss = 0.0 start_time = time.time() for inputs, labels in train_loader: inputs, labels = inputs.to(device), labels.to(device) optimizer.zero_grad() outputs = net(inputs) loss = criterion(outputs, labels) loss.backward() optimizer.step() running_loss += loss.item() end_time = time.time() training_time_per_epoch.append(end_time - start_time) print(f'[Epoch {epoch+1}/{epochs}] Training Loss: {running_loss/len(train_loader):.4f}') # 保存模型状态 state = { 'epoch': epoch, 'model_state_dict': net.state_dict(), 'optimizer_state_dict': optimizer.state_dict() } torch.save(state, model_path) avg_training_time = sum(training_time_per_epoch) / len(training_time_per_epoch) print(f'Average Training Time Per Epoch: {avg_training_time:.2f}s') ``` 在此版本中，利用了 PyTorch 提供的状态字典机制 (`state_dict`) 对模型权重和优化器状态进行了持久化存储[^2]^。这使得程序能够在中断后再继续执行未完成的任务。 --- #### 总结 无论是使用 TensorFlow 或者 PyTorch 构建 CIFAR-10 分类任务中的断点续训功能都非常方便快捷。两者都提供了强大的工具帮助开发者轻松实现这一目标的同时还能有效监控整个训练流程的表现情况如耗时长短等指标变化趋势。