CIFAR10 මත දර්ශන ට්රාන්ස්ෆෝමර් (VIT) පුහුණු කරන්න

13from labml import experiment
14from labml.configs import option
15from labml_nn.experiments.cifar10 import CIFAR10Configs
16from labml_nn.transformers import TransformerConfigs

19class Configs(CIFAR10Configs):

29    transformer: TransformerConfigs

32    patch_size: int = 4

34    n_hidden_classification: int = 2048

36    n_classes: int = 10

39@option(Configs.transformer)
40def _transformer():

44    return TransformerConfigs()

47@option(Configs.model)
48def _vit(c: Configs):

52    from labml_nn.transformers.vit import VisionTransformer, LearnedPositionalEmbeddings, ClassificationHead, \
53        PatchEmbeddings

56    d_model = c.transformer.d_model

58    return VisionTransformer(c.transformer.encoder_layer, c.transformer.n_layers,
59                             PatchEmbeddings(d_model, c.patch_size, 3),
60                             LearnedPositionalEmbeddings(d_model),
61                             ClassificationHead(d_model, c.n_hidden_classification, c.n_classes)).to(c.device)

64def main():

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

68    conf = Configs()

70    experiment.configs(conf, {

72        'optimizer.optimizer': 'Adam',
73        'optimizer.learning_rate': 2.5e-4,

76        'transformer.d_model': 512,

79        'epochs': 32,
80        'train_batch_size': 64,

83        'train_dataset': 'cifar10_train_augmented',

85        'valid_dataset': 'cifar10_valid_no_augment',
86    })

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

90    with experiment.start():
91        conf.run()

95if __name__ == '__main__':
96    main()