【问题标题】:How load torch model for cpu inference如何加载用于 CPU 推理的火炬模型
【发布时间】:2021-06-18 23:00:47
【问题描述】:

无法弄清楚问题所在。即使我将下载指向 cpu,torch 也会尝试使用 cuda 工作:

model = nn.Sequential(
    nn.Conv2d(3,16,5),
    nn.ReLU(),
    nn.MaxPool2d(kernel_size=2, stride=2),
    nn.Conv2d(16,16,5),
    nn.ReLU(),
    nn.MaxPool2d(kernel_size=2, stride=2),
    nn.Conv2d(16,32,5),
    nn.ReLU(),
    nn.MaxPool2d(kernel_size=2, stride=2),
    nn.Conv2d(32,64,5),
    nn.ReLU(),
    nn.MaxPool2d(kernel_size=2, stride=2),
    nn.Conv2d(64,128,5),
    nn.ReLU(),
    nn.MaxPool2d(kernel_size=2, stride=2),
    nn.Flatten(),
    torch.nn.Linear(1152, 256),
    torch.nn.ELU(), torch.nn.Linear(256,4)
            )
model.load_state_dict(torch.load('bbox.pth',map_location=torch.device('cpu')))
ds_trans = transforms.Compose([transforms.ToTensor(),normalize,transforms.Resize((224,224))])
batch = ds_trans(img).unsqueeze(0)
ans = model(batch)

运行结果

AssertionError: Torch not compiled with CUDA enabled

模型是在 GPU 上学习的,GPU 推理运行良好,但我需要在另一台电脑上进行 cpu 推理。

【问题讨论】:

    标签: python pytorch


    【解决方案1】:

    尝试用torch,torchvision重新安装pytorch

    pip install torch===1.5.0 torchvision===0.6.0 -f https://download.pytorch.org/whl/torch_stable.html
    

    【讨论】:

      猜你喜欢
      • 2021-03-01
      • 2021-07-04
      • 2015-09-03
      • 2021-11-14
      • 1970-01-01
      • 2017-08-06
      • 2020-01-23
      • 1970-01-01
      • 1970-01-01
      相关资源
      最近更新 更多