experiment.py

1importtorch2importtorch.nnasnn3fromlabmlimportexperiment4fromlabml.configsimportoption5fromlabml.utils.pytorchimportget\_modules6fromlabml\_helpers.moduleimportModule78fromlabml\_nn.experiments.nlp\_autoregressionimportNLPAutoRegressionConfigs9fromlabml\_nn.hypernetworks.hyper\_lstmimportHyperLSTM10fromlabml\_nn.lstmimportLSTM

13classAutoregressiveModel(Module):

18def\_\_init\_\_(self,n\_vocab:int,d\_model:int,rnn\_model:Module):19super().\_\_init\_\_()

21self.src\_embed=nn.Embedding(n\_vocab,d\_model)22self.lstm=rnn\_model23self.generator=nn.Linear(d\_model,n\_vocab)

25defforward(self,x:torch.Tensor):26x=self.src\_embed(x)

28res,state=self.lstm(x)

30returnself.generator(res),state

33classConfigs(NLPAutoRegressionConfigs):

40model:AutoregressiveModel41rnn\_model:Module4243d\_model:int=51244n\_rhn:int=1645n\_z:int=16

48@option(Configs.model)49defautoregressive\_model(c:Configs):

53m=AutoregressiveModel(c.n\_tokens,c.d\_model,c.rnn\_model)54returnm.to(c.device)

57@option(Configs.rnn\_model)58defhyper\_lstm(c:Configs):59returnHyperLSTM(c.d\_model,c.d\_model,c.n\_rhn,c.n\_z,1)606162@option(Configs.rnn\_model)63deflstm(c:Configs):64returnLSTM(c.d\_model,c.d\_model,1)656667defmain():

69experiment.create(name="hyper\_lstm",comment='')

71conf=Configs()

73experiment.configs(conf,

75{'tokenizer':'character',76'text':'tiny\_shakespeare',77'optimizer.learning\_rate':2.5e-4,78'optimizer.optimizer':'Adam',79'prompt':'It is',80'prompt\_separator':'',8182'rnn\_model':'hyper\_lstm',8384'train\_loader':'shuffled\_train\_loader',85'valid\_loader':'shuffled\_valid\_loader',8687'seq\_len':512,88'epochs':128,89'batch\_size':2,90'inner\_iterations':25})

93experiment.add\_pytorch\_models(get\_modules(conf))

96withexperiment.start():

98conf.run()99100101if\_\_name\_\_=='\_\_main\_\_':102main()