#

CIFAR10 群归一化实验

12import torch.nn as nn
13
14from labml import experiment
15from labml.configs import option
16from labml_helpers.module import Module
17from labml_nn.experiments.cifar10 import CIFAR10Configs
18from labml_nn.normalization.group_norm import GroupNorm

#

用于 CIFAR-10 分类的 VGG 模型

21class Model(Module):

#

26    def __init__(self, groups: int = 32):
27        super().__init__()
28        layers = []

#

RGB 通道

30        in_channels = 3

#

每个区块中每层的通道数

32        for block in [[64, 64], [128, 128], [256, 256, 256], [512, 512, 512], [512, 512, 512]]:

#

卷积、归一化和激活层

34            for channels in block:
35                layers += [nn.Conv2d(in_channels, channels, kernel_size=3, padding=1),
36                           GroupNorm(groups, channels),
37                           nn.ReLU(inplace=True)]
38                in_channels = channels

#

每个区块末端的最大池数

40            layers += [nn.MaxPool2d(kernel_size=2, stride=2)]

#

使用层创建顺序模型

43        self.layers = nn.Sequential(*layers)

#

最后的 logits 层

45        self.fc = nn.Linear(512, 10)

#

47    def forward(self, x):

#

VGG 层

49        x = self.layers(x)

#

修改分类图层的形状

51        x = x.view(x.shape[0], -1)

#

最后的线性层

53        return self.fc(x)

#

56class Configs(CIFAR10Configs):

#

组数

58    groups: int = 16

#

创建模型

61@option(Configs.model)
62def model(c: Configs):

#

66    return Model(c.groups).to(c.device)

#

69def main():

#

创建实验

71    experiment.create(name='cifar10', comment='group norm')

#

创建配置

73    conf = Configs()

#

装载配置

75    experiment.configs(conf, {
76        'optimizer.optimizer': 'Adam',
77        'optimizer.learning_rate': 2.5e-4,
78    })

#

开始实验并运行训练循环

80    with experiment.start():
81        conf.run()

#

85if __name__ == '__main__':
86    main()