pytorch模型权重迁移到megengine

#1

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

1赞
#2

如何从 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 成功之类的，一般问题不大。

#3

补充：在 MegEngine 的 1.9 版本你会神奇地发现 Linear 层的 bias 不需要做 reshape XD