chore: import upstream snapshot with attribution

labml_nn/resnet/experiment.py

@@ -0,0 +1,82 @@
+"""
+---
+title: Train a ResNet on CIFAR 10
+summary: >
+  Train a ResNet on CIFAR 10
+---
+
+# Train a [ResNet](index.html) on CIFAR 10
+"""
+from typing import List, Optional
+
+from torch import nn
+
+from labml import experiment
+from labml.configs import option
+from labml_nn.experiments.cifar10 import CIFAR10Configs
+from labml_nn.resnet import ResNetBase
+
+
+class Configs(CIFAR10Configs):
+    """
+    ## Configurations
+
+    We use [`CIFAR10Configs`](../experiments/cifar10.html) which defines all the
+    dataset related configurations, optimizer, and a training loop.
+    """
+
+    # Number fo blocks for each feature map size
+    n_blocks: List[int] = [3, 3, 3]
+    # Number of channels for each feature map size
+    n_channels: List[int] = [16, 32, 64]
+    # Bottleneck sizes
+    bottlenecks: Optional[List[int]] = None
+    # Kernel size of the initial convolution layer
+    first_kernel_size: int = 3
+
+
+@option(Configs.model)
+def _resnet(c: Configs):
+    """
+    ### Create model
+    """
+    # [ResNet](index.html)
+    base = ResNetBase(c.n_blocks, c.n_channels, c.bottlenecks, img_channels=3, first_kernel_size=c.first_kernel_size)
+    # Linear layer for classification
+    classification = nn.Linear(c.n_channels[-1], 10)
+
+    # Stack them
+    model = nn.Sequential(base, classification)
+    # Move the model to the device
+    return model.to(c.device)
+
+
+def main():
+    # Create experiment
+    experiment.create(name='resnet', comment='cifar10')
+    # Create configurations
+    conf = Configs()
+    # Load configurations
+    experiment.configs(conf, {
+        'bottlenecks': [8, 16, 16],
+        'n_blocks': [6, 6, 6],
+
+        'optimizer.optimizer': 'Adam',
+        'optimizer.learning_rate': 2.5e-4,
+
+        'epochs': 500,
+        'train_batch_size': 256,
+
+        'train_dataset': 'cifar10_train_augmented',
+        'valid_dataset': 'cifar10_valid_no_augment',
+    })
+    # Set model for saving/loading
+    experiment.add_pytorch_models({'model': conf.model})
+    # Start the experiment and run the training loop
+    with experiment.start():
+        conf.run()
+
+
+#
+if __name__ == '__main__':
+    main()