13fromlabmlimportexperiment14fromlabml.configsimportoption15fromlabml\_nn.experiments.cifar10importCIFAR10Configs16fromlabml\_nn.transformersimportTransformerConfigs

19classConfigs(CIFAR10Configs):

29transformer:TransformerConfigs

32patch\_size:int=4

34n\_hidden\_classification:int=2048

36n\_classes:int=10

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

44returnTransformerConfigs()

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

52fromlabml\_nn.transformers.vitimportVisionTransformer,LearnedPositionalEmbeddings,ClassificationHead,\53PatchEmbeddings

56d\_model=c.transformer.d\_model

58returnVisionTransformer(c.transformer.encoder\_layer,c.transformer.n\_layers,59PatchEmbeddings(d\_model,c.patch\_size,3),60LearnedPositionalEmbeddings(d\_model),61ClassificationHead(d\_model,c.n\_hidden\_classification,c.n\_classes)).to(c.device)

64defmain():

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

68conf=Configs()

70experiment.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})

88experiment.add\_pytorch\_models({'model':conf.model})

90withexperiment.start():91conf.run()

95if\_\_name\_\_=='\_\_main\_\_':96main()