thop profile函数遇到nn.DataParallel()时的错误

问题描述：

在神经网络中，常用 From thop import profile 来计算FLOPs和Parameters来作为神经网络模型的评价指标。我在使用该函数时程序报如下错误：RuntimeError: module must have its parameters and buffers on device cuda:0 (device_ids[0]) but found one of them on device: cpu

 解决方式：

我在CSDN上查了半天没有遇到类似的问题，于是上了thop的github官网，企图在issue里找到答案，结果在真给我找到了。issue的网站如下：https://github.com/Lyken17/pytorch-OpCounter/issues/131

我们只需要将源代码：

if torch.cuda.device_count() > 1:
    model = nn.DataParallel(model)

input = torch.randn(1, 8, 224, 224).to(device)
flops, params = profile(model, inputs = (input, ))

修改为：

if torch.cuda.device_count() > 1:
    model = nn.DataParallel(model)

input = torch.randn(1, 8, 224, 224).to(device)
flops, params = profile(model.module, inputs = (input, ))

问题就完美解决了。

另外需要注意的是，如果是用GPU加载模型，则需要在input后面加上to(device)，否则会出现RuntimeError: Input type (torch.FloatTensor) and weight type (torch.cuda.FloatTensor) should be the same报错