AK391
files
d380b77
raw
history blame
932 Bytes
import logging
import torch
from saicinpainting.training.trainers.default import DefaultInpaintingTrainingModule
def get_training_model_class(kind):
if kind == 'default':
return DefaultInpaintingTrainingModule
raise ValueError(f'Unknown trainer module {kind}')
def make_training_model(config):
kind = config.training_model.kind
kwargs = dict(config.training_model)
kwargs.pop('kind')
kwargs['use_ddp'] = config.trainer.kwargs.get('accelerator', None) == 'ddp'
logging.info(f'Make training model {kind}')
cls = get_training_model_class(kind)
return cls(config, **kwargs)
def load_checkpoint(train_config, path, map_location='cuda', strict=True):
model: torch.nn.Module = make_training_model(train_config)
state = torch.load(path, map_location=map_location)
model.load_state_dict(state['state_dict'], strict=strict)
model.on_load_checkpoint(state)
return model