File size: 950 Bytes
762f6ed b7c3fbd 762f6ed | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 | # -*- coding: utf-8 -*-
#Author: Lart Pang (https://github.com/lartpang)
import os
import torch
def save_weight(save_path, model):
print(f"Saving weight '{save_path}'")
if isinstance(model, dict):
model_state = model
else:
model_state = model.module.state_dict() if hasattr(model, "module") else model.state_dict()
torch.save(model_state, save_path)
print(f"Saved weight '{save_path}' " f"(only contain the net's weight)")
def load_weight(load_path, model, *, strict=True, skip_unmatched_shape=False):
assert os.path.exists(load_path), load_path
model_params = model.state_dict()
for k, v in torch.load(load_path, map_location="cpu").items():
if k.endswith("module."):
k = k[7:]
if skip_unmatched_shape and k in model_params and v.shape != model_params[k].shape:
continue
model_params[k] = v
model.load_state_dict(model_params, strict=strict)
|