-
Notifications
You must be signed in to change notification settings - Fork 1
/
utils.py
44 lines (40 loc) · 1.56 KB
/
utils.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
import os
import os.path
import torch
import sys
class saveData():
def __init__(self, args):
self.args = args
self.save_dir = os.path.join(args.saveDir, args.name)
#self.save_dir = os.path.join(args.saveDir, args.load)
if not os.path.exists(self.save_dir):
os.makedirs(self.save_dir)
self.save_dir_model = os.path.join(self.save_dir, 'model')
if not os.path.exists(self.save_dir_model):
os.makedirs(self.save_dir_model)
if os.path.exists(self.save_dir + '/log.txt'):
self.logFile = open(self.save_dir + '/log.txt', 'a')
else:
self.logFile = open(self.save_dir + '/log.txt', 'w')
def save_model(self, model, epoch):
torch.save(
model.state_dict(),
self.save_dir_model + '/model_latest.pt')
torch.save(
model.state_dict(),
self.save_dir_model + '/model_' + str(epoch) + '.pt')
torch.save(
model,
self.save_dir_model + '/model_obj.pt')
torch.save(
epoch,
self.save_dir_model + '/last_epoch.pt')
def save_log(self, log):
sys.stdout.flush()
self.logFile.write(log + '\n')
self.logFile.flush()
def load_model(self, model):
model.load_state_dict(torch.load(self.save_dir_model + '/model_latest.pt'))
last_epoch = torch.load(self.save_dir_model + '/last_epoch.pt')
print("load mode_status frmo {}/model_latest.pt, epoch: {}".format(self.save_dir_model, last_epoch))
return model, last_epoch