Question

我对如何重新加载具有不同批次大小的pytorch模型有疑问。在训练中，我的批处理大小为64，但据推断，我希望批处理大小为1（一次喂入数据）。这是我用来保存和恢复模型的代码：

torch.save(agent.qnetwork_local.state_dict(), './ckpt/checkpoint.pth')
saved_model = QNetwork(state_size=37, action_size=4, seed=0)
saved_model.load_state_dict(torch.load('./ckpt/checkpoint.pth'))

运行推理模型时出现此错误：

RuntimeError: size mismatch, m1: [37 x 1], m2: [37 x 64] at /Users/soumith/code/builder/wheel/pytorch-src/aten/src/TH/generic/THTensorMath.cpp:2070

此错误表示模型的输入必须为37x64，其中37是数据尺寸，而64是训练批次大小。但是测试输入为37x1，这意味着数据尺寸为37，批量大小为1。

在重装pytorch模型中是否有针对不同批次大小的解决方案？非常感谢。

Answer 1

我最终在DataLoader中使用batch_size=1做到了

import torch
import pandas as pd
from torch.utils.data.dataloader import DataLoader

df = pd.read_csv('data.csv')
df = df.values

# Use CustomDataset class for your data
inference_dataset = CustomDataset(x=df[:1, 0:2])

inference_dataloader = DataLoader(inference_dataset, batch_size=1, shuffle=False, num_workers=4)

# 
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load('./model/model'))
model.eval()

for i, x in enumerate(inference_dataloader):
    x = x.float()
    y_pred = model(x)
    print(y_pred.value)

Answer 2

建立模型时，可以使用-1动态表示批次大小。例如，下面是前阶段代码

def forward(self, x):
     x = self.conv1(x)
     x = self.layer1(x)
     x = self.layer2(x)
     x = self.avgpool(x)
     x = x.view(-1, 37)
 #instead using x.view(64,37) 
     x = self.fc(x)

希望它可以帮助您

pytorch恢复具有不同批次大小的模型

2 个答案: