-
Notifications
You must be signed in to change notification settings - Fork 1
/
utils.py
75 lines (61 loc) · 2.72 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
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
import numpy as np
import torch
from torch import optim
import torch.nn as nn
import argparse
import yaml
import os
import shutil
def save_checkpoint(state, is_best, checkpoint_dir):
"""Saves model and training parameters at '{checkpoint_dir}/last_checkpoint.pytorch'.
If is_best==True saves '{checkpoint_dir}/best_checkpoint.pytorch' as well.
Args:
state (dict): contains model's state_dict, optimizer's state_dict, epoch
and best evaluation metric value so far
is_best (bool): if True state contains the best model seen so far
checkpoint_dir (string): directory where the checkpoint are to be saved
"""
if not os.path.exists(checkpoint_dir):
os.mkdir(checkpoint_dir)
last_file_path = os.path.join(checkpoint_dir, 'last_checkpoint.pytorch')
torch.save(state, last_file_path)
if is_best:
best_file_path = os.path.join(checkpoint_dir, 'best_checkpoint.pytorch')
shutil.copyfile(last_file_path, best_file_path)
def load_checkpoint(checkpoint_path, model, optimizer=None,
model_key='model_state_dict', optimizer_key='optimizer_state_dict'):
"""Loads model and training parameters from a given checkpoint_path
If optimizer is provided, loads optimizer's state_dict of as well.
Args:
checkpoint_path (string): path to the checkpoint to be loaded
model (torch.nn.Module): model into which the parameters are to be copied
optimizer (torch.optim.Optimizer) optional: optimizer instance into
which the parameters are to be copied
Returns:
state
"""
if not os.path.exists(checkpoint_path):
raise IOError(f"Checkpoint '{checkpoint_path}' does not exist")
state = torch.load(checkpoint_path, map_location='cpu')
model.load_state_dict(state[model_key])
if optimizer is not None:
optimizer.load_state_dict(state[optimizer_key])
def load_config():
parser = argparse.ArgumentParser(description='Resnet')
parser.add_argument('--config', type=str, help='Path to the YAML config file', required=True)
args = parser.parse_args()
config = yaml.safe_load(open(args.config, 'r'))
# Get a device to train on
device_str = config.get('device', None)
if device_str is not None:
logger.info(f"Device specified in config: '{device_str}'")
if device_str.startswith('cuda') and not torch.cuda.is_available():
logger.warning('CUDA not available, using CPU')
device_str = 'cpu'
else:
device_str = "cuda:0" if torch.cuda.is_available() else 'cpu'
device = torch.device(device_str)
config['device'] = device
return config
def _load_config_yaml(config_file):
return yaml.safe_load(open(config_file, 'r'))