pytorch模型权重迁移到megengine

能有提供一个教程或者示例,如何把pytorch模型的权重迁移到megengine模型上?

1赞

如何从 PyTorch 迁移 weight 到 MegEngine 中?

首先确保两边的 Module 结构是类似的,在 PyTorch 里保存 state_dict,然后在 MegEngine 中将这个 weight 载入即可。

# 在 PyTorch 中保存权重
import pickle
with open('torch-weight.pkl', 'wb') as f:
    states = net.state_dict()
    weights = {k: v.numpy() for k, v in states.items()}
    pickle.dump(weights, f)
# 在 MegEngine 中读取权重
import pickle
with open('torch-weight.pkl', 'rb') as f:
    w = pickle.load(f)
weights = {}
for k, v in w.items():
    if k.endswith('bias') and v.ndim == 1:
        v = v.reshape(1, -1, 1, 1)
    weights[k] = v

net.load_state_dict(weights, strict=False)

在这个过程中可能会遇到一些 warning,可能是一些统计量没 load 成功之类的,一般问题不大。