How to make an ensemble model for classification with pytorch using trained models?

Question

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)

Answer 1

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).