Open xingjchen opened 3 years ago
I think I get the right solution now.
We just need to replace the '_fc' layer:
class AuxNet(torch.nn.Module):
def __init__(self, input_features):
super(AuxNet, self).__init__()
self.fc1 = torch.nn.Sequential( torch.nn.Linear(input_features, 512))
self.fc2 = torch.nn.Sequential( torch.nn.Linear(input_features, 128))
def forward(self, x):
y1 = self.fc1(x)
y2 = self.fc2(x)
return y1,y2
effi_model = EfficientNet.from_pretrained('efficientnet-b7')
feature = effi_model._fc.in_features
effi_model._fc = AuxNet(feature)
X = torch.rand(8, 3,224,224)
effi_model(X)
And this kind of solution can be applied to all similar cases.
Hi, I have a problem with multi-task learning when using EfficientNet.
I want to load the pre-trained model and remove the raw FC layers, then add my own 2 new FC layers for 2 different tasks.
However, it always shows the following bug:
File "/home/user/miniconda3/envs/pytorch/lib/python3.8/site-packages/torch/nn/modules/module.py", line 889, in _call_impl result = self.forward(*input, **kwargs) TypeError: forward() takes 1 positional argument but 2 were given
and here is my code:
Is there any solution? There is the same issue proposed here https://stackoverflow.com/questions/62954999/flattening-efficientnet-model
Thank you very much!