|
"""Tools to save/restore model from checkpoints.""" |
|
|
|
import argparse |
|
import shutil |
|
import sys |
|
import os |
|
import re |
|
import json |
|
import time |
|
|
|
import torch |
|
|
|
CHECKPOINT_PATTERN = re.compile('^model_checkpoint-(\d+)$') |
|
|
|
|
|
class ArgsDict(dict): |
|
|
|
def __init__(self, **kwargs): |
|
super(ArgsDict, self).__init__() |
|
for key, value in kwargs.items(): |
|
self[key] = value |
|
self.__dict__ = self |
|
|
|
|
|
def load_checkpoint(item_dict, model_dir, map_location=None, step=None): |
|
""" item_dict: {"model": model, "opt1": opt1, ...}""" |
|
path = os.path.join(model_dir, 'model_checkpoint') |
|
if step is not None: |
|
path += '-{:08d}'.format(step) |
|
if os.path.exists(path): |
|
print("Loading model from %s" % path) |
|
checkpoint = torch.load(path, map_location=map_location) |
|
|
|
old_state_dict = item_dict["model"].state_dict() |
|
for key in old_state_dict.keys(): |
|
if key not in checkpoint['model']: |
|
checkpoint['model'][key] = old_state_dict[key] |
|
|
|
for item_name in item_dict: |
|
item_dict[item_name].load_state_dict(checkpoint[item_name]) |
|
return checkpoint.get('step', 0) |
|
return 0 |
|
|
|
|
|
def load_and_map_checkpoint(model, model_dir, remap): |
|
path = os.path.join(model_dir, 'model_checkpoint') |
|
print("Loading parameters %s from %s" % (remap.keys(), model_dir)) |
|
checkpoint = torch.load(path) |
|
new_state_dict = model.state_dict() |
|
for name, value in remap.items(): |
|
|
|
new_state_dict[name] = checkpoint['model'][value] |
|
model.load_state_dict(new_state_dict) |
|
|
|
|
|
def save_checkpoint(items, step, model_dir, ignore=[], |
|
keep_every_n=10000000): |
|
if not os.path.exists(model_dir): |
|
os.makedirs(model_dir) |
|
path_without_step = os.path.join(model_dir, 'model_checkpoint') |
|
step_padded = format(step, '08d') |
|
state_dict = items["model"].state_dict() |
|
if ignore: |
|
for key in state_dict.keys(): |
|
for item in ignore: |
|
if key.startswith(item): |
|
state_dict.pop(key) |
|
path_with_step = '{}-{}'.format(path_without_step, step_padded) |
|
|
|
saved_dic = {} |
|
for key in items: |
|
saved_dic[key] = items[key].state_dict() |
|
torch.save({**saved_dic, "step": step}, path_with_step) |
|
|
|
try: |
|
os.unlink(path_without_step) |
|
except FileNotFoundError: |
|
pass |
|
try: |
|
os.symlink(os.path.basename(path_with_step), path_without_step) |
|
except OSError: |
|
shutil.copy2(path_with_step, path_without_step) |
|
|
|
|
|
if keep_every_n is not None: |
|
all_checkpoints = [] |
|
for name in os.listdir(model_dir): |
|
m = CHECKPOINT_PATTERN.match(name) |
|
if m is None or name == os.path.basename(path_with_step): |
|
continue |
|
checkpoint_step = int(m.group(1)) |
|
all_checkpoints.append((checkpoint_step, name)) |
|
all_checkpoints.sort() |
|
|
|
last_step = float('-inf') |
|
for checkpoint_step, name in all_checkpoints: |
|
if checkpoint_step - last_step >= keep_every_n: |
|
last_step = checkpoint_step |
|
continue |
|
os.unlink(os.path.join(model_dir, name)) |
|
|
|
|
|
class Saver(object): |
|
"""Class to manage save and restore for the model and optimizer.""" |
|
|
|
def __init__(self, items, keep_every_n=None): |
|
assert type(items) == dict |
|
assert "model" in items |
|
self._items = items |
|
self._keep_every_n = keep_every_n |
|
|
|
def restore(self, model_dir, map_location=None, |
|
step=None, item_keys=["model", "optimizer"]): |
|
"""Restores model and optimizer from given directory. |
|
Specify what shoud be restored |
|
|
|
Returns: |
|
Last training step for the model restored. |
|
""" |
|
items2restore = { k: self._items[k] for k in item_keys} |
|
last_step = load_checkpoint( |
|
items2restore, model_dir, map_location, step) |
|
return last_step |
|
|
|
def save(self, model_dir, step): |
|
"""Saves model and optimizer to given directory. |
|
|
|
Args: |
|
model_dir: Model directory to save. |
|
step: Current training step. |
|
""" |
|
save_checkpoint(self._items, step, model_dir, |
|
keep_every_n=self._keep_every_n) |
|
|
|
def restore_part(self, other_model_dir, remap): |
|
"""Restores part of the model from other directory. |
|
|
|
Useful to initialize part of the model with another pretrained model. |
|
|
|
Args: |
|
other_model_dir: Model directory to load from. |
|
remap: dict, remapping current parameters to the other model's. |
|
""" |
|
load_and_map_checkpoint(self._items["model"], other_model_dir, remap) |
|
|
