def save_state(
filename, args, model_state_dict, criterion, optimizer, lr_scheduler,
num_updates, optim_history=None, extra_state=None,
):
from fairseq import utils
if optim_history is None:
optim_history = []
if extra_state is None:
extra_state = {}
state_dict = {
'args': args,
'model': model_state_dict if model_state_dict else {},
'optimizer_history': optim_history + [
{
'criterion_name': criterion.__class__.__name__,
'optimizer_name': optimizer.__class__.__name__,
'lr_scheduler_state': lr_scheduler.state_dict(),
'num_updates': num_updates,
}
],
'extra_state': extra_state,
}
if utils.has_parameters(criterion):
state_dict['criterion'] = criterion.state_dict()
if not args.no_save_optimizer_state:
state_dict['last_optimizer_state'] = convert_state_dict_type(optimizer.state_dict())
try:
from fairseq.fb_pathmgr import fb_pathmgr
with fb_pathmgr.open(filename, "wb") as f:
torch_persistent_save(state_dict, f)
except (ModuleNotFoundError, ImportError):
# if path manager not found, continue with local file.
torch_persistent_save(state_dict, filename)
def save_state(
filename,
args,
model_state_dict,
criterion,
optimizer,
lr_scheduler,
num_updates,
optim_history=None,
extra_state=None,
):
from fairseq import utils
if optim_history is None:
optim_history = []
if extra_state is None:
extra_state = {}
state_dict = {
'args':
args,
'model':
model_state_dict if model_state_dict else {},
'optimizer_history':
optim_history + [{
'criterion_name': criterion.__class__.__name__,
'optimizer_name': optimizer.__class__.__name__,
'lr_scheduler_state': lr_scheduler.state_dict(),
'num_updates': num_updates,
}],
'extra_state':
extra_state,
}
if utils.has_parameters(criterion):
state_dict['criterion'] = criterion.state_dict()
if not args.no_save_optimizer_state:
state_dict['last_optimizer_state'] = convert_state_dict_type(
optimizer.state_dict())
#if args.criterion == 'semantic_similarity_loss': # TODO : refactoring required
# criterion.rewarder.save_ckpt(os.path.join("./semsim/trained_models/", filename.split('/')[-1]))
try:
from fairseq.fb_pathmgr import fb_pathmgr
with fb_pathmgr.open(filename, "wb") as f:
torch_persistent_save(state_dict, f)
except (ModuleNotFoundError, ImportError):
# if path manager not found, continue with local file.
torch_persistent_save(state_dict, filename)
def save_state(
filename, args, model_state_dict, criterion, optimizer, lr_scheduler,
num_updates, optim_history=None, extra_state=None,
):
from fairseq import utils
if optim_history is None:
optim_history = []
if extra_state is None:
extra_state = {}
state_dict = {
'args': args,
'model': model_state_dict if model_state_dict else {},
'optimizer_history': optim_history + [
{
'criterion_name': criterion.__class__.__name__,
'optimizer_name': optimizer.__class__.__name__,
'lr_scheduler_state': lr_scheduler.state_dict(),
'num_updates': num_updates,
}
],
'extra_state': extra_state,
}
if utils.has_parameters(criterion):
state_dict['criterion'] = criterion.state_dict()
if not args.no_save_optimizer_state:
state_dict['last_optimizer_state'] = convert_state_dict_type(optimizer.state_dict())
xla = not getattr(args, 'use_gpu', True)
try:
from fairseq.fb_pathmgr import fb_pathmgr
with fb_pathmgr.open(filename, "wb") as f:
torch_persistent_save(state_dict, f, xla=xla)
except (ModuleNotFoundError, ImportError):
if filename.startswith(gcsfs.CLOUD_STORAGE_PREFIX):
# tpu-comment: Saving to GCS directly is problematic at the moment.
tmp_filename = os.path.join('/tmp', os.path.basename(filename))
torch_persistent_save(
state_dict, tmp_filename, xla=xla, global_master=True
)
if xm.is_master_ordinal(local=False):
with open(tmp_filename, 'rb') as fd:
gcsfs.write(filename, fd.read())
os.remove(tmp_filename)
else:
# continue with local file.
torch_persistent_save(state_dict, filename, xla=xla)
def load_checkpoint_to_cpu(path, arg_overrides=None):
"""Loads a checkpoint to CPU (with upgrading for backward compatibility)."""
try:
from fairseq.fb_pathmgr import fb_pathmgr
with fb_pathmgr.open(path, "rb") as f:
state = torch.load(
f, map_location=lambda s, l: default_restore_location(s, 'cpu'),
)
except (ModuleNotFoundError, ImportError):
# if path manager not found, continue with local file.
state = torch.load(
path, map_location=lambda s, l: default_restore_location(s, 'cpu'),
)
args = state['args']
if arg_overrides is not None:
for arg_name, arg_val in arg_overrides.items():
setattr(args, arg_name, arg_val)
state = _upgrade_state_dict(state)
return state
