| import torch |
| from .ema import ExponentialMovingAverage |
|
|
| def load_model_weights(model, ckpt_path, use_ema=True, device='cuda:0'): |
| """ |
| Load weights of a model from a checkpoint file. |
| |
| Args: |
| model (torch.nn.Module): The model to load weights into. |
| ckpt_path (str): Path to the checkpoint file. |
| use_ema (bool): Whether to use Exponential Moving Average (EMA) weights if available. |
| """ |
| checkpoint = torch.load(ckpt_path,map_location={'cuda:0': str(device)}) |
| total_iter = checkpoint.get('total_it', 0) |
| |
| if "model_ema" in checkpoint and use_ema: |
| ema_key = next(iter(checkpoint["model_ema"])) |
| if ('module' in ema_key) or ('n_averaged' in ema_key): |
| model = ExponentialMovingAverage(model, decay=1.0) |
| |
| model.load_state_dict(checkpoint["model_ema"], strict=True) |
| if ('module' in ema_key) or ('n_averaged' in ema_key): |
| model = model.module |
| print(f'\nLoading EMA module model from {ckpt_path} with {total_iter} iterations') |
| else: |
| print(f'\nLoading EMA model from {ckpt_path} with {total_iter} iterations') |
| else: |
| model.load_state_dict(checkpoint['encoder'], strict=True) |
| print(f'\nLoading model from {ckpt_path} with {total_iter} iterations') |
|
|
| return total_iter |
