def train():
train_model = Director(batch_size=args.batch_size,
gpu=args.gpu,
num_workers=args.num_workers,
hidden_size=args.hidden_size,
model_dir=args.model_dir)
if args.resume:
train_model.load_ckpts()
train_model.train(epochs=args.num_epochs,
lr=args.lr,
log_every=args.log_every)
def test():
test_model = Director(batch_size=args.batch_size,
gpu=args.gpu,
num_workers=args.num_workers,
hidden_size=args.hidden_size,
model_dir=args.model_dir)
epoch, _ = test_model.load_ckpts()
test_model.evaluate()
test_model.visualize(epoch, 'test')
