时间戳

首页 > 解决方案 > 具有不同批量大小的pytorch恢复模型

python - 具有不同批量大小的pytorch恢复模型

问题描述

我有一个关于如何重新加载具有不同批量大小的 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模型中的不同批量大小有什么解决方案吗？非常感谢。

标签： pythonneural-networkdeep-learningpytorch

解决方案

我最终设法batch_size=1在DataLoader中使用

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)

推荐阅读