-
Notifications
You must be signed in to change notification settings - Fork 33
/
Copy pathmain.py
31 lines (25 loc) · 838 Bytes
/
main.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
import torch
import utility
import model as model_module
import loss as loss_module
from option import args
from trainer import Trainer
import dataloader.dataset as dataset
torch.manual_seed(args.seed)
checkpoint = utility.checkpoint(args)
def main():
if checkpoint.ok:
loader = dataset.myData(args)
model = model_module.Model(args, checkpoint)
print('Total params: %.2fM' % (sum(p.numel() for p in model.parameters())/1000000.0))
loss = loss_module.Loss(args, checkpoint) if not args.test_only else None
trainer = Trainer(args, loader, model, loss, checkpoint)
if args.test_only:
trainer.test()
else:
for i in range(args.epochs):
trainer.train()
trainer.test()
checkpoint.done()
if __name__ == '__main__':
main()