-
Notifications
You must be signed in to change notification settings - Fork 2
/
loader.py
51 lines (44 loc) · 1.88 KB
/
loader.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
51
# coding: utf-8
import torch
import os
import glob
import numpy as np
def load_previous_model(encoder, decoder, checkpoint_dir, model_prefix):
f_list = glob.glob(os.path.join(checkpoint_dir, model_prefix) + '-*.pth')
start_epoch = 1
if len(f_list) >= 1:
epoch_list = [int(i.split('-')[-1].split('.')[0]) for i in f_list]
last_checkpoint = f_list[np.argmax(epoch_list)]
if os.path.exists(last_checkpoint):
print('load from {}'.format(last_checkpoint))
model_state_dict = torch.load(last_checkpoint, map_location=lambda storage, loc: storage)
encoder.load_state_dict(model_state_dict['encoder'])
decoder.load_state_dict(model_state_dict['decoder'])
start_epoch = np.max(epoch_list)
return encoder, decoder, start_epoch
def save_model(encoder, decoder, checkpoint_dir, model_prefix, epoch, max_keep = 5):
"""
params:
encoder: the encoder used for learning
decoder: the decoder used for learning
checkpoint_dir: directory where the model will save
model_prefix:
epoch: epoch number
max_keep: How much model should be stored
if not os.path.exists(checkpoint_dir):
os.makedirs(checkpoint_dir)
f_list = glob.glob(os.path.join(checkpoint_dir, model_prefix) + '-*.pth')
if len(f_list) >=max_keep + 2:
# this step using for delete the more than 5 and litter one
epoch_list = [int(i.split('-')[-1].split('.')[0]) for i in f_list]
to_delete = [f_list[i] for i in np.argsort(epoch_list)[-max_keep:]]
for f in to_delete:
os.remove(f)
name = model_prefix + '-{}.pth'.format(epoch)
file_path = os.path.join(checkpoint_dir, name)
model_dict = {
'encoder': encoder.state_dict(),
'decoder': decoder.state_dict()
}
torch.save(model_dict, file_path)
print("Saved model.")