Train a small model on CIFAR 10

This trains a small model on CIFAR 10 to test how much distillation benefits.

13import torch.nn as nn
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


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 SmallModel(CIFAR10VGGModel):

Create a convolution layer and the activations

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

Convolution layer

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

Batch normalization

46            BatchNorm(out_channels, track_running_stats=False),

ReLU activation

48            nn.ReLU(inplace=True),
49        )
51    def __init__(self):

Create a model with given convolution sizes (channels)

53        super().__init__([[32, 32], [64, 64], [128], [128], [128]])

Create model

57def _small_model(c: Configs):
61    return SmallModel().to(c.device)
64def main():

Create experiment

66    experiment.create(name='cifar10', comment='small model')

Create configurations

68    conf = Configs()

Load configurations

70    experiment.configs(conf, {
71        'optimizer.optimizer': 'Adam',
72        'optimizer.learning_rate': 2.5e-4,
73    })

Set model for saving/loading

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

Print number of parameters in the model

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

Start the experiment and run the training loop

79    with experiment.start():

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