#

Train a large model on CIFAR 10

This trains a large model on CIFAR 10 for distillation.

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

#

Configurations

We use CIFAR10Configs which defines all the dataset related configurations, optimizer, and a training loop.

21class Configs(CIFAR10Configs):

#

28    pass

#

VGG style model for CIFAR-10 classification

This derives from the generic VGG style architecture.

31class LargeModel(CIFAR10VGGModel):

#

Create a convolution layer and the activations

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

#

42        return nn.Sequential(

#

Dropout

44            nn.Dropout(0.1),

#

Convolution layer

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

#

Batch normalization

48            BatchNorm(out_channels, track_running_stats=False),

#

ReLU activation

50            nn.ReLU(inplace=True),
51        )

#

53    def __init__(self):

#

Create a model with given convolution sizes (channels)

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

#

Create model

58@option(Configs.model)
59def _large_model(c: Configs):

#

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

#

66def main():

#

Create experiment

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

#

Create configurations

70    conf = Configs()

#

Load configurations

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

#

Set model for saving/loading

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

#

Print number of parameters in the model

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

#

Start the experiment and run the training loop

83    with experiment.start():
84        conf.run()

#

88if __name__ == '__main__':
89    main()