-
Notifications
You must be signed in to change notification settings - Fork 5
/
run.py
50 lines (35 loc) · 1.47 KB
/
run.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
import fire
import utils
from torch import optim
from models.model import Model
from settings.hparam import hparam as hp
class Runner:
IMPLEMENTED_MODELS = ['cbhg', 'mgru']
def train(self, model, checkpoint='', is_cuda=True, is_multi_gpu=True, logdir='', savedir=''):
if model not in self.IMPLEMENTED_MODELS:
raise NotImplementedError('%s model is not implemented !' % model)
mode = 'train'
logger = utils.get_logger(mode)
# initialize hyperparameters
hp.set_hparam_yaml(mode)
logger.info('Setup mode as %s, model : %s' % (mode, model))
# get network
network = utils.get_networks(model, checkpoint, is_cuda, is_multi_gpu)
# setup dataset
train_dataloader = Model.data_loader(mode='train')
test_dataloader = Model.data_loader(mode='test')
# setup optimizer:
parameters = network.parameters()
logger.info(network)
# TODO: Scheduled LR
lr = getattr(hp, mode).lr
optimizer = optim.Adam([p for p in parameters if p.requires_grad], lr=lr)
# pass model, loss, optimizer and dataset to the trainer
# get trainer
trainer = utils.get_trainer()(network, optimizer, train_dataloader, test_dataloader, is_cuda, logdir, savedir)
# train!
trainer.run(hp.train.num_epochs)
def eval(self):
raise NotImplementedError('Evaluation mode is not implemented!')
if __name__ == '__main__':
fire.Fire(Runner)