fastaiv2 到 pytorch for torchserver

fastaiv2 to pytorch for torchserver

我通常使用 fastai(v2 或 v1)进行快速原型制作。现在我想部署我的一个模型,用 fastai 训练,到 torchserver。

假设我们有一个像这样的简单模型:

learn = cnn_learner(data, 
                models.resnet34, 
                metrics=[accuracy, error_rate, score])
# after the training 
torch.save(learn.model.state_dict(), "./test1.pth")
state = torch.load("./test1.pth")
model_torch_rep = models.resnet34()
model_torch_rep.load_state_dict(state)

我试过很多不同的东西,结果都一样

RuntimeError                              Traceback (most recent call last)
<ipython-input-284-e4dbdce23d43> in <module>
----> 1 model_torch_rep.load_state_dict(state);

/opt/conda/lib/python3.6/site-packages/torch/nn/modules/module.py in load_state_dict(self, state_dict, strict)
    837         if len(error_msgs) > 0:
    838             raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format(
--> 839                                self.__class__.__name__, "\n\t".join(error_msgs)))
    840         return _IncompatibleKeys(missing_keys, unexpected_keys)
    841 

RuntimeError: Error(s) in loading state_dict for ResNet:
    Missing key(s) in state_dict: "conv1.weight", "bn1.weight", "bn1.bias", "bn1.running_mean", "bn1.running_var", "layer1.0.conv1.weight", "layer1.0.bn1.weight"

这发生在 fastai 1.0.6 或 fastai 2.3.1 + pytorch 1.8.1 ...

刚刚弄明白了。

出于某种原因,您保存 state_dict 的方式添加了一个字符串“module”。已加载 state_dict 中的每个键。 (我想这是因为您没有使用 FastAI 的 Learner class 来保存模型)。

只需删除“模块”。来自状态字典的子字符串,你一切都好。

learn = cnn_learner(data, 
                    models.resnet34, 
                    metrics=[accuracy, error_rate, score])

# after the training 
torch.save(learn.model.state_dict(), "./test1.pth")
state = torch.load("./test1.pth")

# fix dict keys 
new_state = OrderedDict([(k.partition('module.')[2], v) for k, v in state.items()])

model_torch_rep = models.resnet34()
model_torch_rep.load_state_dict(new_state)