Pytorch模型重量类型转换

时间:2018-03-16 05:49:03

标签: pytorch

我试图从文件中对FlowNet2-C模型加载进行推断。 但是,我遇到了一些数据类型问题。我该如何解决?

$ python main.py

Initializing Datasets
  [0.000s] Loading checkpoint '/notebooks/data/model/FlowNet2-C_checkpoint.pth.tar'
  [1.293s] Loaded checkpoint '/notebooks/data/model/FlowNet2-C_checkpoint.pth.tar' (at epoch 0)
(1L, 6L, 384L, 512L)
<class 'torch.autograd.variable.Variable'>
  [1.642s] Operation failed


Traceback (most recent call last):
  File "main.py", line 102, in <module>
    main()
  File "main.py", line 98, in main
    summary(input_size, model)
  File "main.py", line 61, in summary
    model(x)
  File "/usr/local/lib/python2.7/dist-packages/torch/nn/modules/module.py", line 357, in __call__
    result = self.forward(*input, **kwargs)
  File "/notebooks/data/vinet/FlowNetC.py", line 75, in forward
    out_conv1a = self.conv1(x1)
  File "/usr/local/lib/python2.7/dist-packages/torch/nn/modules/module.py", line 357, in __call__
    result = self.forward(*input, **kwargs)
  File "/usr/local/lib/python2.7/dist-packages/torch/nn/modules/container.py", line 67, in forward
    input = module(input)
  File "/usr/local/lib/python2.7/dist-packages/torch/nn/modules/module.py", line 357, in __call__
    result = self.forward(*input, **kwargs)
  File "/usr/local/lib/python2.7/dist-packages/torch/nn/modules/conv.py", line 282, in forward
    self.padding, self.dilation, self.groups)
  File "/usr/local/lib/python2.7/dist-packages/torch/nn/functional.py", line 90, in conv2d
    return f(input, weight, bias)
RuntimeError: Input type (CUDAFloatTensor) and weight type (CPUFloatTensor) should be the same

1 个答案:

答案 0 :(得分:5)

也许这是因为您对模型的model和输入x具有不同的数据类型。似乎模型参数已移至GPU,但您的输入x在GPU上。

您可以尝试在第94行之后使用model.cuda(),这会将模型放在GPU上。然后错误就会消失。