CIFAR10 හි විශාල ආකෘතියක් පුහුණු කරන්න

15import torch.nn as nn
16
17from labml import experiment, logger
18from labml.configs import option
19from labml_nn.experiments.cifar10 import CIFAR10Configs, CIFAR10VGGModel
20from labml_nn.normalization.batch_norm import BatchNorm

23class Configs(CIFAR10Configs):

30    pass

33class LargeModel(CIFAR10VGGModel):

40    def conv_block(self, in_channels, out_channels) -> nn.Module:

44        return nn.Sequential(

46            nn.Dropout(0.1),

48            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),

50            BatchNorm(out_channels, track_running_stats=False),

52            nn.ReLU(inplace=True),
53        )

55    def __init__(self):

57        super().__init__([[64, 64], [128, 128], [256, 256, 256], [512, 512, 512], [512, 512, 512]])

60@option(Configs.model)
61def _large_model(c: Configs):

65    return LargeModel().to(c.device)

68def main():

70    experiment.create(name='cifar10', comment='large model')

72    conf = Configs()

74    experiment.configs(conf, {
75        'optimizer.optimizer': 'Adam',
76        'optimizer.learning_rate': 2.5e-4,
77        'is_save_models': True,
78        'epochs': 20,
79    })

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

83    logger.inspect(params=(sum(p.numel() for p in conf.model.parameters() if p.requires_grad)))

85    with experiment.start():
86        conf.run()

90if __name__ == '__main__':
91    main()