#

在 CIFAR 10 上训练 R esNet

12from typing import List, Optional
13
14from torch import nn
15
16from labml import experiment
17from labml.configs import option
18from labml_nn.experiments.cifar10 import CIFAR10Configs
19from labml_nn.resnet import ResNetBase

#

配置

我们使用CIFAR10Configs 它来定义所有与数据集相关的配置、优化器和训练循环。

22class Configs(CIFAR10Configs):

#

每个要素地图大小的区块数

31    n_blocks: List[int] = [3, 3, 3]

#

每个要素映射大小的通道数

33    n_channels: List[int] = [16, 32, 64]

#

瓶颈大小

35    bottlenecks: Optional[List[int]] = None

#

初始卷积层的内核大小

37    first_kernel_size: int = 3

#

创建模型

40@option(Configs.model)
41def _resnet(c: Configs):

#

ResNet

46    base = ResNetBase(c.n_blocks, c.n_channels, c.bottlenecks, img_channels=3, first_kernel_size=c.first_kernel_size)

#

用于分类的线性层

48    classification = nn.Linear(c.n_channels[-1], 10)

#

堆叠它们

51    model = nn.Sequential(base, classification)

#

将模型移到设备上

53    return model.to(c.device)

#

56def main():

#

创建实验

58    experiment.create(name='resnet', comment='cifar10')

#

创建配置

60    conf = Configs()

#

装载配置

62    experiment.configs(conf, {
63        'bottlenecks': [8, 16, 16],
64        'n_blocks': [6, 6, 6],
65
66        'optimizer.optimizer': 'Adam',
67        'optimizer.learning_rate': 2.5e-4,
68
69        'epochs': 500,
70        'train_batch_size': 256,
71
72        'train_dataset': 'cifar10_train_augmented',
73        'valid_dataset': 'cifar10_valid_no_augment',
74    })

#

设置保存/加载的模型

76    experiment.add_pytorch_models({'model': conf.model})

#

开始实验并运行训练循环

78    with experiment.start():
79        conf.run()

#

83if __name__ == '__main__':
84    main()