map2map/map2map/state.py
Yin Li d2840f01b0 Add --adv-start as epoch to start adversarial training
This is similar to the deprecated --adv-delay, but specify
absolute epoch (--adv-delay counts from start_epoch in resumed training)
2020-02-13 16:00:50 -05:00

16 lines
503 B
Python

import warnings
import sys
from pprint import pformat
def load_model_state_dict(module, state_dict, strict=True):
bad_keys = module.load_state_dict(state_dict, strict)
if len(bad_keys.missing_keys) > 0:
warnings.warn('Missing keys in state_dict:\n{}'.format(
pformat(bad_keys.missing_keys)))
if len(bad_keys.unexpected_keys) > 0:
warnings.warn('Unexpected keys in state_dict:\n{}'.format(
pformat(bad_keys.unexpected_keys)))
sys.stderr.flush()