-
Notifications
You must be signed in to change notification settings - Fork 1
/
Data.py
57 lines (43 loc) · 1.72 KB
/
Data.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
52
53
54
55
56
57
import torch
import os
def save_model(model, optimizer, loss_func, scheduler,
epoch, train_loss, test_loss,
BATCH_SIZE, LR, WEIGHT_DECAY,
INSTRUMENT, ):
torch.save({
'model_name': model.__class__.__name__,
'model_state_dict': model.state_dict(),
'optimizer': optimizer,
'optimizer_state_dict': optimizer.state_dict(),
'loss_func': loss_func,
'scheduler': scheduler,
'epoch': epoch,
'train_loss': train_loss,
'test_loss': test_loss,
'BATCH_SIZE': BATCH_SIZE,
'LR': LR,
'WEIGHT_DECAY': WEIGHT_DECAY,
'INSTRUMENT' : INSTRUMENT,
},
f'model/{model.__class__.__name__}.pkl')
def load_model(path):
checkpoint = torch.load(path)
model = eval(f'architecture.{checkpoint["model_name"]}().to(device)')
model.load_state_dict(checkpoint['model_state_dict'])
return model
def load(path):
checkpoint = torch.load(path)
model = eval(f'architecture.{checkpoint["model_name"]}().to(device)')
model.load_state_dict(checkpoint['model_state_dict'])
optimizer = checkpoint['optimizer']
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
loss_func = checkpoint['loss_func']
scheduler = checkpoint['scheduler']
return model, optimizer, loss_func, scheduler
def load_hyperparam(path):
checkpoint = torch.load(path)
SIZE = checkpoint['SIZE']
BATCH_SIZE = checkpoint['BATCH_SIZE']
LR = checkpoint['LR']
WEIGHT_DECAY = checkpoint['WEIGHT_DECAY']
return SIZE, BATCH_SIZE, LR, WEIGHT_DECAY