PyTorch 工具、最佳实践与风格指南

这不是 PyTorch 的官方风格指南。本文档总结了我们使用 PyTorch 框架进行深度学习一年多以来的最佳实践。请注意，我们分享的经验主要来自研究和初创公司的视角。

这是一个开放项目，欢迎其他贡献者编辑和改进本文档。

您将在这份文档中找到三个主要部分。首先是对 Python 最佳实践的快速回顾，接着是一些使用 PyTorch 的技巧和建议。最后，我们将分享一些使用其他框架的见解和经验，这些帮助我们整体上提升了工作流程。

更新 2020年12月20日

• 添加了一个在 cifar10 数据集上训练模型的完整示例
• 增加了使用 VS Code 和远程扩展的设置指南

更新 2019年4月30日

在收到大量积极反馈后，我还添加了我们在 Lightly 项目中常用的一些构建模块的摘要： 您将找到用于（自注意力机制、基于 VGG 的感知损失、谱归一化、自适应实例归一化等）的构建模块。
损失函数、层及其他构建模块的代码片段

我们推荐使用 Python 3.6+

根据我们的经验，我们建议使用 Python 3.6 或更高版本，因为以下特性对编写整洁简洁的代码非常有帮助：

• Python 3.6 起支持类型注解。
• Python 3.6 起支持 f-string。

Python 风格指南回顾

我们尽量遵循 Google 的 Python 风格指南。 请参考 Google 提供的、文档完善的 Python 代码风格指南。

在此我们提供最常用规则的摘要：

命名规范

摘自 3.16.4

类型	规范	示例
包与模块	小写加下划线	from prefetch_generator import BackgroundGenerator
类	首字母大写	class DataLoader
常量	全大写加下划线	BATCH_SIZE=16
实例	小写加下划线	dataset = Dataset
方法与函数	小写加下划线()	def visualize_tensor()
变量	小写加下划线	background_color='Blue'

IDEs

代码编辑器

一般来说，我们推荐使用像 Visual Studio Code 或 PyCharm 这样的 IDE。虽然 VS Code 是一个相对轻量级的编辑器，却提供了语法高亮和自动补全功能，而 PyCharm 则拥有许多用于处理远程集群的高级特性。 VS Code 凭借其快速增长的扩展生态系统，已经变得非常强大。

使用远程机器设置 Visual Studio Code

请确保已安装以下扩展：

• Python（语法检查、自动补全、语法高亮、代码格式化）
• Remote - SSH（用于连接远程机器）

1. 按照此处的指南操作：https://code.visualstudio.com/docs/remote/remote-overview

设置 PyCharm 以连接远程机器

1. 登录到您的远程机器（AWS、Google 等）。
2. 创建一个新的文件夹和一个新的虚拟环境。
3. 在 PyCharm（专业版）的项目设置中配置远程解释器。
4. 配置远程 Python 解释器（指向 AWS、Google 等上的虚拟环境路径）。
5. 配置本地代码与远程机器之间的映射关系。

如果设置正确，这将使您能够做到以下几点：

• 在本地计算机上编写代码（笔记本电脑、台式机），无论您身在何处（离线或在线）。
• 将本地代码同步到远程机器。
• 额外的包会自动安装在远程机器上。
• 您无需在本地机器上存储任何数据集。
• 在远程机器上运行和调试代码，就像在本地机器上运行一样。

Jupyter Notebook 与 Python 脚本

一般来说，我们建议在初步探索或尝试新模型和代码时使用 Jupyter Notebook。 一旦您需要在更大的数据集上训练模型，并且更注重可重复性时，则应使用 Python 脚本。

我们推荐的工作流程：

1. 从 Jupyter Notebook 开始。
2. 探索数据和模型。
3. 在 Notebook 的单元格中构建类和方法。
4. 将代码迁移到 Python 脚本中。
5. 在服务器上进行训练或部署。

Jupyter Notebook	Python 脚本
+ 适合探索	+ 可以长时间运行任务而不中断
+ 方便调试	+ 易于通过 git 跟踪更改
- 容易变成超大文件	- 调试通常意味着重新运行整个脚本
- 可能会被中断（不适合长时间训练）
- 容易出错并变得混乱

库

常用库：

名称	描述	用途
torch	用于神经网络的基础框架	创建张量、构建网络并使用反向传播进行训练
torchvision	PyTorch 计算机视觉模块	图像数据预处理、增强和后处理
Pillow (PIL)	Python 图像处理库	加载和保存图像
Numpy	用于科学计算的 Python 库	数据的预处理和后处理
prefetch_generator	后台处理库	在计算过程中后台加载下一批数据
tqdm	进度条	每个 epoch 训练过程中的进度显示
torchinfo	为 PyTorch 提供类似 Keras 的模型摘要	显示网络结构、各层参数及尺寸
torch.utils.tensorboard	PyTorch 内置 TensorBoard	记录实验并在 TensorBoard 中可视化

文件组织

不要把所有的层和模型都放在同一个文件里。最佳实践是将最终的网络单独放在一个文件中（例如 networks.py），并将层、损失函数和操作分别放在各自的文件中（layers.py、losses.py、ops.py）。完成的模型（由一个或多个网络组成）应在以其命名的文件中被引用（如 yolov3.py、DCGAN.py）。

主程序以及训练和测试脚本只应从包含模型名称的文件中导入内容。

在 PyTorch 中构建神经网络

我们建议将网络拆分为较小的可重用组件。一个网络是一个 nn.Module，由操作或其他 nn.Module 作为构建块组成。损失函数也是 nn.Module，因此可以直接集成到网络中。

继承自 nn.Module 的类必须有一个 forward 方法，用于实现相应层或操作的前向传播。

可以使用 self.net(input) 将一个 nn.module 应用于输入数据。这实际上就是调用该对象的 call() 方法，将输入通过模块传递。

output = self.net(input)

PyTorch 中的一个简单网络

对于只有一个输入和一个输出的简单网络，可以使用以下模式：

class ConvBlock(nn.Module):
    def __init__(self):
        super(ConvBlock, self).__init__()
        self.block = nn.Sequential(
            nn.Conv2d(...), 
            nn.ReLU(), 
            nn.BatchNorm2d(...)
        )  
    
    def forward(self, x):
        return self.block(x)

class SimpleNetwork(nn.Module):
    def __init__(self, num_resnet_blocks=6):
        super(SimpleNetwork, self).__init__()
        # 在这里添加各个层
        layers = [ConvBlock(...)]
        for i in range(num_resnet_blocks):
            layers += [ResBlock(...)]
        self.net = nn.Sequential(*layers)
    
    def forward(self, x):
        return self.net(x)

请注意以下几点：

• 我们重复使用简单的、可重复的构建块，例如 ConvBlock，它由相同的循环模式（卷积、激活、归一化）组成，并将其放入一个单独的 nn.Module 中。
• 我们构建一个所需层的列表，最后使用 nn.Sequential() 将其转换为模型。我们在列表对象前使用 * 操作符来解包列表。
• 在前向传播中，我们只需将输入通过模型即可。

在 PyTorch 中具有跳跃连接的网络

class ResnetBlock(nn.Module):
    def __init__(self, dim, padding_type, norm_layer, use_dropout, use_bias):
        super(ResnetBlock, self).__init__()
        self.conv_block = self.build_conv_block(...)

    def build_conv_block(self, ...):
        conv_block = []

        conv_block += [nn.Conv2d(...),
                       norm_layer(...),
                       nn.ReLU()]
        if use_dropout:
            conv_block += [nn.Dropout(...)]
            
        conv_block += [nn.Conv2d(...),
                       norm_layer(...)]

        return nn.Sequential(*conv_block)

    def forward(self, x):
        out = x + self.conv_block(x)
        return out

在这里，ResNet 块的跳跃连接直接在前向传播中实现。PyTorch 允许在前向传播过程中进行动态操作。

在 PyTorch 中具有多个输出的网络

对于需要多个输出的网络，例如使用预训练的 VGG 网络构建感知损失时，可以使用以下模式：

class Vgg19(nn.Module):
  def __init__(self, requires_grad=False):
    super(Vgg19, self).__init__()
    vgg_pretrained_features = models.vgg19(pretrained=True).features
    self.slice1 = torch.nn.Sequential()
    self.slice2 = torch.nn.Sequential()
    self.slice3 = torch.nn.Sequential()

    for x in range(7):
        self.slice1.add_module(str(x), vgg_pretrained_features[x])
    for x in range(7, 21):
        self.slice2.add_module(str(x), vgg_pretrained_features[x])
    for x in range(21, 30):
        self.slice3.add_module(str(x), vgg_pretrained_features[x])
    if not requires_grad:
        for param in self.parameters():
            param.requires_grad = False

  def forward(self, x):
    h_relu1 = self.slice1(x)
    h_relu2 = self.slice2(h_relu1)        
    h_relu3 = self.slice3(h_relu2)        
    out = [h_relu1, h_relu2, h_relu3]
    return out

请注意以下几点：

• 我们使用了 torchvision 提供的预训练模型。
• 我们将网络拆分为三个切片。每个切片由预训练模型中的层组成。
• 我们通过设置 requires_grad = False 来“冻结”网络。
• 我们返回一个包含三个切片输出的列表。

自定义损失函数

即使 PyTorch 已经提供了许多标准损失函数，有时仍然可能需要创建自己的损失函数。为此，可以创建一个单独的文件 losses.py，并扩展 nn.Module 类来创建自定义损失函数：

class CustomLoss(nn.Module):
    
    def __init__(self):
        super(CustomLoss,self).__init__()
        
    def forward(self,x,y):
        loss = torch.mean((x - y)**2)
        return loss

训练模型的推荐代码结构

完整的示例可以在本仓库的 cifar10-example 文件夹中找到。

请注意，我们使用了以下模式：

• 我们使用 prefetch_generator 中的 BackgroundGenerator 在后台加载下一个批次 更多信息请参见此问题。
• 我们使用 tqdm 监控训练进度并显示“计算效率”。这有助于我们找到数据加载管道中的瓶颈。

# 导入语句
import torch
import torch.nn as nn
from torch.utils import data
...

# 设置标志/种子
torch.backends.cudnn.benchmark = True
np.random.seed(1)
torch.manual_seed(1)
torch.cuda.manual_seed(1)
...

# 从主代码开始
if __name__ == '__main__':
    # 使用 argparse 处理实验的额外参数
    parser = argparse.ArgumentParser(description="训练一个用于...的网络")
    ...
    opt = parser.parse_args() 
    
    # 添加数据集相关代码（我们通常使用训练集和验证/测试集）
    data_transforms = transforms.Compose([
        transforms.Resize((opt.img_size, opt.img_size)),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
    ])
    
    train_dataset = datasets.ImageFolder(
        root=os.path.join(opt.path_to_data, "train"),
        transform=data_transforms)
    train_data_loader = data.DataLoader(train_dataset, ...)
    
    test_dataset = datasets.ImageFolder(
        root=os.path.join(opt.path_to_data, "test"),
        transform=data_transforms)
    test_data_loader = data.DataLoader(test_dataset ...)
    ...
    
    # 实例化网络（已从 *networks.py* 中导入）
    net = MyNetwork(...)
    ...
    
    # 创建损失函数（PyTorch 中的 criterion）
    criterion_L1 = torch.nn.L1Loss()
    ...
    
    # 如果在 GPU 上运行且希望使用 CUDA，则将模型移动到 GPU
    use_cuda = torch.cuda.is_available()
    if use_cuda:
        net = net.cuda()
        ...
    
    # 创建优化器
    optim = torch.optim.Adam(net.parameters(), lr=opt.lr)
    ...
    
    # 如果需要或希望，加载检查点
    start_n_iter = 0
    start_epoch = 0
    if opt.resume:
        ckpt = load_checkpoint(opt.path_to_checkpoint) # 自定义方法，用于加载上次的检查点
        net.load_state_dict(ckpt['net'])
        start_epoch = ckpt['epoch']
        start_n_iter = ckpt['n_iter']
        optim.load_state_dict(ckpt['optim'])
        print("已恢复上次的检查点")
        ...
        
    # 如果希望在多个 GPU 上运行实验，则将模型转移到这些 GPU 上
    net = torch.nn.DataParallel(net)
    ...
    
    # 通常我们会使用 tensorboardX 来跟踪实验过程
    writer = SummaryWriter(...)
    
    # 现在开始主循环
    n_iter = start_n_iter
    for epoch in range(start_epoch, opt.epochs):
        # 将模型设置为训练模式
        net.train()
        ...
        
        # 使用 prefetch_generator 和 tqdm 遍历数据
        pbar = tqdm(enumerate(BackgroundGenerator(train_data_loader, ...)),
                    total=len(train_data_loader))
        start_time = time.time()
        
        # 遍历数据集的循环
        for i, data in pbar:
            # 数据准备
            img, label = data
            if use_cuda:
                img = img.cuda()
                label = label.cuda()
            ...
            
            # 使用 tqdm 跟踪数据准备时间和计算时间是非常好的做法，这样可以发现数据加载器中的问题
            prepare_time = start_time-time.time()
            
            # 前向传播和反向传播
            optim.zero_grad()
            ...
            loss.backward()
            optim.step()
            ...
            
            # 更新 tensorboardX
            writer.add_scalar(..., n_iter)
            ...
            
            # 计算计算时间和 *计算效率*
            process_time = start_time-time.time()-prepare_time
            pbar.set_description("计算效率: {:.2f}, 第 {}/{} 个 epoch:".format(
                process_time/(process_time+prepare_time), epoch, opt.epochs))
            start_time = time.time()
            
        # 可以每隔 x 个 epoch 进行一次测试
        if epoch % x == x-1:
            # 将模型切换到评估模式
            net.eval()
            ...
            # 进行一些测试
            pbar = tqdm(enumerate(BackgroundGenerator(test_data_loader, ...)),
                    total=len(test_data_loader)) 
            for i, data in pbar:
                ...
                
            # 如果需要，保存检查点
            ...

在 PyTorch 中使用多 GPU 进行训练

在 PyTorch 中有两种不同的模式可用于多 GPU 训练。 根据我们的经验，这两种模式都是可行的。不过第一种方式更简洁、代码量更少。第二种方式由于 GPU 之间的通信较少，性能上略有优势。关于这两种方法的讨论，请参阅 PyTorch 官方论坛

将每个网络的批量输入拆分

最常见的方式是将所有 网络 的批次简单地分配到各个 GPU 上。

例如，如果一个模型在单个 GPU 上以 64 的批量大小运行，那么在两个 GPU 上运行时，每个 GPU 的批量大小则为 32。这可以通过用 nn.DataParallel(model) 包装模型来自动完成。

将所有网络打包成一个“超级”网络，并拆分输入批次

这种模式相对较少使用。NVIDIA 的 pix2pixHD 实现 就是一个采用这种方法的仓库。

应该做的和不应该做的

避免在 nn.Module 的前向传播方法中使用 NumPy 代码

NumPy 是在 CPU 上运行的，速度比 PyTorch 代码慢。由于 PyTorch 在设计时就考虑了与 NumPy 的相似性，因此大多数 NumPy 函数都已经在 PyTorch 中得到了支持。

将数据加载器与主代码分离

数据加载流程应独立于主训练代码。PyTorch 使用后台工作线程来更高效地加载数据，同时不会干扰主训练过程。

不要在每一步都记录结果

通常我们的模型会训练数千步。因此，只需每隔几步记录一次损失和其他结果即可，这样可以减少开销。尤其是在训练过程中，保存中间结果图像可能会非常耗时。

使用命令行参数

使用命令行参数来设置代码执行期间的参数（如批量大小、学习率等）非常方便。一种简单的方法是打印从 parse_args 接收到的参数字典，以便跟踪实验参数：

...
# 将参数保存到 config.txt 文件
opt = parser.parse_args()
with open("config.txt", "w") as f:
    f.write(opt.__str__())
...

尽可能使用 .detach() 从计算图中释放张量

PyTorch 会跟踪所有涉及张量的操作以实现自动微分。使用 .detach() 可以避免记录不必要的操作。

打印标量张量时使用 .item()

可以直接打印变量，但建议使用 variable.detach() 或 variable.item()。在 PyTorch 0.4 之前的版本中，需要使用 .data 来访问变量的张量。

在 nn.Module 上使用 call 方法而不是 forward 方法

这两种方法并不完全相同，正如其中一个议题所指出的那样 这里：

output = self.net.forward(input)

# 它们不相等！
output = self.net(input)

常见问题解答

1. 如何确保我的实验可复现？

我们建议在代码开头设置以下随机种子：

np.random.seed(1)
torch.manual_seed(1)
torch.cuda.manual_seed(1)

2. 如何进一步提升训练和推理速度？

在 Nvidia GPU 上，您可以在代码开头添加以下一行。这将使 CUDA 后端在首次执行时优化您的计算图。但请注意，如果更改了网络的输入或输出张量大小，每次发生改变时都会重新优化计算图。这可能导致运行速度非常慢，并出现内存不足错误。只有当您的输入和输出始终保持相同形状时，才应启用此标志。通常情况下，这会带来约 20% 的性能提升。

torch.backends.cudnn.benchmark = True

3. 使用 tqdm + prefetch_generator 模式时，计算效率的理想值是多少？

这取决于所使用的机器、预处理流水线以及网络规模。在配备 1080Ti GPU 和 SSD 的设备上，我们观察到接近 1.0 的计算效率，这是理想情况。如果使用较浅（较小）的网络或速度较慢的硬盘，该数值可能会降至 0.1–0.2 左右，具体取决于您的配置。

4. 如果显存不足，如何实现批量大小大于 1？ 在 PyTorch 中，我们可以非常容易地实现虚拟批量大小。只需阻止优化器更新参数，并在 batch_size 次循环中累计梯度即可。

...
# 在主循环中
out = net(input)
loss = criterion(out, label)
# 只调用 backward 累计梯度，但不在此处执行 step
loss.backward() 
total_loss += loss.item() / batch_size
if n_iter % batch_size == batch_size-1:
    # 此处使用虚拟批量大小执行优化步骤
    optim.step()
    optim.zero_grad()
    print('总损失：', total_loss)
    total_loss = 0.0
...

5. 如何在训练过程中调整学习率？

我们可以直接通过实例化的优化器访问学习率，如下所示：

...
for param_group in optim.param_groups:
    old_lr = param_group['lr']
    new_lr = old_lr * 0.1
    param_group['lr'] = new_lr
    print('学习率从 {} 调整为 {}'.format(old_lr, new_lr))
...

6. 如何在训练过程中将预训练模型用作损失函数（不进行反向传播）？

如果您希望使用 VGG 等预训练模型来计算损失，但不对其进行训练（例如，在风格迁移、GAN 或自编码器中使用的感知损失），可以采用以下模式：

...
# 实例化模型
pretrained_VGG = VGG19(...)

# 禁用梯度（防止训练）
for p in pretrained_VGG.parameters():  # 将 requires_grad 重置为 False
    p.requires_grad = False
...
# 不必使用 no_grad() 上下文管理器，直接运行模型即可
# VGG 模型不会计算任何梯度
out_real = pretrained_VGG(input_a)
out_fake = pretrained_VGG(input_b)
loss = any_criterion(out_real, out_fake)
...

7. 为什么在 PyTorch 中要使用 .train() 和 .eval() 方法？

这些方法用于将 BatchNorm2d 或 Dropout2d 等层从训练模式切换到推理模式。所有继承自 nn.Module 的模块都具有一个名为 isTraining 的属性。.eval() 和 .train() 只是简单地将该属性设置为 True/False。有关此方法实现的更多信息，请参阅 PyTorch 源码中的模块部分。

8. 我的模型在推理时占用大量内存。如何在 PyTorch 中正确地进行推理？ 请确保在代码执行过程中不计算和存储任何梯度。您可以使用以下模式来保证这一点：

with torch.no_grad():
    # 在此处运行模型
    out_tensor = net(in_tensor)

9. 如何对预训练模型进行微调？

在 PyTorch 中，您可以冻结某些层，从而阻止它们在优化步骤中被更新。

# 您可以通过以下方式冻结整个模块：
for p in pretrained_VGG.parameters():  # 将 requires_grad 重置为 False
    p.requires_grad = False

10. 何时使用 Variable(...)？

自 PyTorch 0.4 版本起，Variable 和 Tensor 已经合并。我们不再需要显式创建 Variable 对象。

11. 使用 C++ 编写的 PyTorch 是否比 Python 版本更快？ C++ 版本大约快 10%
12. TorchScript / JIT 能否加速我的代码？

待办……

13. 使用 cudnn.benchmark=True 的 PyTorch 代码是否更快？ 根据我们的经验，可以提升约 20% 的速度。不过，首次运行模型时，构建优化后的计算图需要较长时间。在某些情况下（前向传播中存在循环、输入形状不固定、前向传播中包含 if/else 等），此标志可能会导致“内存不足”或其他错误。
14. 如何使用多块 GPU 进行训练？

待办……

15. PyTorch 中的 .detach() 是如何工作的？ 它会将张量从计算图中分离出来。一个很好的示意图可以在这里找到：http://www.bnikolic.co.uk/blog/pytorch-detach.html

您喜欢这个仓库吗？

请提供反馈，帮助我们改进这份风格指南！您可以开一个 issue，或者通过创建 pull request 提出修改建议。

如果您喜欢这个仓库，别忘了查看我们提供的其他框架：

• Lightly - 用于自监督学习的计算机视觉框架

PyTorch Styleguide 快速上手指南

本指南基于 pytorch-styleguide 项目整理，旨在帮助开发者遵循最佳实践，构建清晰、可维护且高效的 PyTorch 深度学习代码。

环境准备

在开始之前，请确保满足以下系统和依赖要求：

• 操作系统: Linux, macOS 或 Windows
• Python 版本: 推荐 Python 3.6+
  • 原因：支持类型提示（Type Hinting）和 f-strings，有助于编写更整洁的代码。
• 核心框架: PyTorch 及 torchvision
• 常用辅助库:
  • Pillow: 图像处理
  • Numpy: 科学计算
  • prefetch_generator: 后台数据加载加速
  • tqdm: 进度条显示
  • torchinfo: 模型结构摘要
• 开发工具 (可选但推荐):
  • VS Code: 需安装 Python 和 Remote - SSH 扩展（适用于远程服务器开发）。
  • PyCharm Professional: 支持配置远程解释器和代码映射。

安装步骤

建议使用虚拟环境进行隔离。以下是基于 pip 的安装命令（国内用户可使用清华源加速）：

# 创建并激活虚拟环境
python -m venv venv
source venv/bin/activate  # Windows 使用: venv\Scripts\activate

# 安装 PyTorch (请以官网最新命令为准，此处为示例)
# 国内加速源推荐：-i https://pypi.tuna.tsinghua.edu.cn/simple
pip install torch torchvision torchaudio -i https://pypi.tuna.tsinghua.edu.cn/simple

# 安装项目推荐的辅助工具库
pip install pillow numpy prefetch_generator tqdm torchinfo -i https://pypi.tuna.tsinghua.edu.cn/simple

基本使用

本项目主要提供代码规范、架构建议和实用代码片段，而非单一的可执行命令行工具。以下是核心用法的快速演示：

1. 项目文件组织建议

不要将所有代码堆砌在一个文件中。推荐结构如下：

• networks.py: 定义最终的网络模型
• layers.py: 定义可复用的层模块
• losses.py: 定义自定义损失函数
• ops.py: 定义其他操作
• train.py / test.py: 主训练/测试脚本（仅导入模型文件）

2. 构建模块化网络

将网络拆分为可复用的小块（继承 nn.Module），并使用 nn.Sequential 组装。

import torch
import torch.nn as nn

# 定义可复用的基础模块
class ConvBlock(nn.Module):
    def __init__(self):
        super(ConvBlock, self).__init__()
        self.block = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=3, padding=1), 
            nn.ReLU(), 
            nn.BatchNorm2d(64)
        )  
    
    def forward(self, x):
        return self.block(x)

# 组装主网络
class SimpleNetwork(nn.Module):
    def __init__(self, num_blocks=6):
        super(SimpleNetwork, self).__init__()
        layers = [ConvBlock()]
        # 动态添加层
        for i in range(num_blocks):
            layers += [ConvBlock()] 
        self.net = nn.Sequential(*layers)
    
    def forward(self, x):
        return self.net(x)

# 实例化并使用
model = SimpleNetwork()
input_tensor = torch.randn(1, 3, 64, 64)
output = model(input_tensor)
print(output.shape)

3. 实现自定义损失函数

通过继承 nn.Module 轻松创建自定义 Loss。

class CustomLoss(nn.Module):
    def __init__(self):
        super(CustomLoss, self).__init__()
        
    def forward(self, x, y):
        # 示例：均方误差
        loss = torch.mean((x - y)**2)
        return loss

criterion = CustomLoss()
# loss = criterion(prediction, target)

4. 推荐的工作流

1. 探索阶段: 使用 Jupyter Notebook 进行数据探索和原型验证。
2. 重构阶段: 将验证过的代码类和方法迁移至 Python 脚本 (.py)。
3. 训练阶段: 在服务器上使用 Python 脚本进行长时间训练，配合 tqdm 监控进度，使用 prefetch_generator 优化数据加载。

# 训练循环中的最佳实践示例
from tqdm import tqdm
from prefetch_generator import BackgroundGenerator

# 设置随机种子以保证可复现性
torch.backends.cudnn.benchmark = True
np.random.seed(1)
torch.manual_seed(1)

# 使用 BackgroundGenerator 加速数据加载
for batch in BackgroundGenerator(data_loader, max_prefetch=3):
    # 训练逻辑...
    pass

提示: 完整的 CIFAR-10 训练示例请参考该项目仓库中的 cifar10-example 文件夹。