0
$\begingroup$

I am trying to make an ensemble model composed of two pre-trained models, using torch, in order to classify an image.

Below is some code, based on this post.

import timm
import torch
from torch.nn import functional as F

num_classes = 100

model1 = timm.create_model("efficientnet_b0", num_classes=num_classes)
checkpoint1 = torch.load(checkpoint_path1)
model1.load_state_dict(checkpoint1["model"])

model2 = timm.create_model("efficientnet_b2", num_classes=num_classes)
checkpoint2 = torch.load(checkpoint_path2)
model2.load_state_dict(checkpoint2["model"])

class EnsembleModel(nn.Module):
    def __init__(self, modelA, modelB, num_features):
        super().__init__()
        self.modelA = modelA
        self.modelB = modelB
        self.classifier = nn.Linear(2*num_features, num_features)
        
    def forward(self, x):
        x1 = self.modelA(x)
        x2 = self.modelB(x)
        x = torch.cat((x1, x2), dim=1)
        x = self.classifier(F.relu(x))
        return x

model = EnsembleModel(model1, model2, num_classes)

The two models have been pre-trained on a GPU (cuda), and when I run a prediction from EnsembleModel, I get this error:

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cpu and cuda:0! (when checking arugment for argument mat1 in method wrapper_addmm)

from this line:

---> 15         x = self.classifier(F.relu(x))

Honestly, I'm not even sure why the post suggested using a classifier, and combining them with a relu.

What is the best way to combine two models like this?

Here is more of the stack trace if that is useful:

/tmp/ipykernel_33/4277289726.py in forward(self, x)
     13         x2 = self.modelB(x)
     14         x = torch.cat((x1, x2), dim=1)
---> 15         x = self.classifier(F.relu(x))
     16         return x

/opt/conda/lib/python3.7/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
   1049         if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1050                 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1051             return forward_call(*input, **kwargs)
   1052         # Do not call functions when jit is used
   1053         full_backward_hooks, non_full_backward_hooks = [], []

/opt/conda/lib/python3.7/site-packages/torch/nn/modules/linear.py in forward(self, input)
     94 
     95     def forward(self, input: Tensor) -> Tensor:
---> 96         return F.linear(input, self.weight, self.bias)
     97 
     98     def extra_repr(self) -> str:

/opt/conda/lib/python3.7/site-packages/torch/nn/functional.py in linear(input, weight, bias)
   1845     if has_torch_function_variadic(input, weight):
   1846         return handle_torch_function(linear, (input, weight), input, weight, bias=bias)
-> 1847     return torch._C._nn.linear(input, weight, bias)
   1848 
   1849 

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cpu and cuda:0! (when checking arugment for argument mat1 in method wrapper_addmm)
$\endgroup$

1 Answer 1

0
$\begingroup$

From your code snippet, my best guess is that the cpu and cuda:0 error does not mean your 2 models but input data x and self.classifier. You can check their locations by printing them. Then you can move the one on the CPU to GPU using .to(device).

$\endgroup$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.