Spaces:
Runtime error
Runtime error
import json | |
import logging | |
import os | |
import sys | |
import torch | |
MATPLOTLIB_FLAG = False | |
logging.basicConfig(stream=sys.stdout, level=logging.DEBUG) | |
logger = logging | |
def load_checkpoint(checkpoint_path, model, optimizer=None): | |
assert os.path.isfile(checkpoint_path) | |
checkpoint_dict = torch.load(checkpoint_path, map_location='cpu') | |
iteration = checkpoint_dict['iteration'] | |
learning_rate = checkpoint_dict['learning_rate'] | |
if optimizer is not None: | |
optimizer.load_state_dict(checkpoint_dict['optimizer']) | |
saved_state_dict = checkpoint_dict['model'] | |
if hasattr(model, 'module'): | |
state_dict = model.module.state_dict() | |
else: | |
state_dict = model.state_dict() | |
new_state_dict = {} | |
for k, v in state_dict.items(): | |
try: | |
new_state_dict[k] = saved_state_dict[k] | |
except Exception as e: | |
logger.info(e) | |
logger.info("%s is not in the checkpoint" % k) | |
new_state_dict[k] = v | |
if hasattr(model, 'module'): | |
model.module.load_state_dict(new_state_dict) | |
else: | |
model.load_state_dict(new_state_dict) | |
logger.info("Loaded checkpoint '{}' (iteration {})".format( | |
checkpoint_path, iteration)) | |
return model, optimizer, learning_rate, iteration | |
def save_checkpoint(model, optimizer, learning_rate, iteration, checkpoint_path): | |
logger.info("Saving model and optimizer state at iteration {} to {}".format( | |
iteration, checkpoint_path)) | |
if hasattr(model, 'module'): | |
state_dict = model.module.state_dict() | |
else: | |
state_dict = model.state_dict() | |
torch.save({'model': state_dict, | |
'iteration': iteration, | |
'optimizer': optimizer.state_dict(), | |
'learning_rate': learning_rate}, checkpoint_path) | |
def get_hparams_from_file(config_path): | |
with open(config_path, "r", encoding="utf-8") as f: | |
data = f.read() | |
config = json.loads(data) | |
hparams = HParams(**config) | |
return hparams | |
class HParams: | |
def __init__(self, **kwargs): | |
for k, v in kwargs.items(): | |
if type(v) == dict: | |
v = HParams(**v) | |
self[k] = v | |
def keys(self): | |
return self.__dict__.keys() | |
def items(self): | |
return self.__dict__.items() | |
def values(self): | |
return self.__dict__.values() | |
def __len__(self): | |
return len(self.__dict__) | |
def __getitem__(self, key): | |
return getattr(self, key) | |
def __setitem__(self, key, value): | |
return setattr(self, key, value) | |
def __contains__(self, key): | |
return key in self.__dict__ | |
def __repr__(self): | |
return self.__dict__.__repr__() | |