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