#

グループ正規化のための 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)

#

最終ロジットレイヤー

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()