Problem using pre-trained model as a layer

I am trying use a pre-trained resnet model as a layer in my ConvNet class.

class ConvNet(nn.Module):

 def __init__(self,body,C):
     self.body = body
     self.head = nn.Linear(1000,C)
 def forward(x):
     return self.head(self.body(x))

import torchvision.models as models
body = models.resnet18(pretrained=True)
conv_net = ConvNet(body,5)

And I get this error:

in init(self, body, C)
8 def init(self,body,C):
----> 9 self.body = body
10 self.head = nn.Linear(1000,C)

~/anaconda3/envs/fastai/lib/python3.6/site-packages/torch/nn/modules/ in setattr(self, name, value)
422 if modules is None:
423 raise AttributeError(
–> 424 “cannot assign module before Module.init() call”)
425 remove_from(self.dict, self._parameters, self._buffers)
426 modules[name] = value

AttributeError: cannot assign module before Module.init() call

I am not sure what the problem is, I thought the body would be initialized if I used <pretrained=True>.

It looks like you need to call the __init__ of the parent class:

class ConvNet(nn.Module):
    def __init__(self, body, C):
        super(ConvNet, self).__init__()
        self.body = ...

As a small side note, if you want to use the pretrained model as a fixed feature extractor, the usual way would be to replace the last linear layer with your own instead of stacking it on top of the logits. But maybe that’s exactly your use case. :wink:

Oh, that was silly. Thanks, works now.