如何加载 Torch model 进行 CPU 推理

Question

Can't figure out what the problem is. Even if I point the download to cpu, torch tries to use cuda to work:

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)

result of run

AssertionError: Torch not compiled with CUDA enabled

Model was learned on GPU, and GPU inference is working well, but I need cpu inference on another pc.

Answer 1

Try to Reinstall the pytorch with torch,torchvision

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