EfficientNet-PyTorch
EfficientNet-PyTorch copied to clipboard
lukemelas
Question regarding last layers
Hello, I got a question: what does passing num_classes actually do? add another extra linear layer at the end with corresponding classes? What happens if I don't pass it? How can I get the same network as when I pass num_classes =8 with an additional dropout layer? Best, Finn
Not sure if that's what you want, but you can probably do the following
class MyEfficientNet(nn.Module):
def __init__(self):
super().__init__()
# EfficientNet
self.network = EfficientNet.from_pretrained("efficientnet-b0")
# Replace last layer
self.network._fc = nn.Sequential(nn.Linear(self.network._fc.in_features, 512),
nn.ReLU(),
nn.Dropout(0.25),
nn.Linear(512, 128),
nn.ReLU(),
nn.Dropout(0.50),
nn.Linear(128,classes))
def forward(self, x):
out = self.network(x)
return out
model = MyEfficientNet()
And add as much layers you want. as the num of classes, efficientnet has 1000 outputs because of ImageNet, so changing that just make easier for you to train your custom dataset without having to rewrite the last layer.
Not sure if that's what you want, but you can probably do the following
class MyEfficientNet(nn.Module): def __init__(self): super().__init__() # EfficientNet self.network = EfficientNet.from_pretrained("efficientnet-b0") # Replace last layer self.network._fc = nn.Sequential(nn.Linear(self.network._fc.in_features, 512), nn.ReLU(), nn.Dropout(0.25), nn.Linear(512, 128), nn.ReLU(), nn.Dropout(0.50), nn.Linear(128,classes)) def forward(self, x): out = self.network(x) return out model = MyEfficientNet()And add as much layers you want. as the num of classes, efficientnet has 1000 outputs because of ImageNet, so changing that just make easier for you to train your custom dataset without having to rewrite the last layer.
I tried you code, but starts to get GPU memory error ,even for light b0, before i could use b2 for train
"CUDA out of memory. Tried to allocate 22.00 MiB (GPU 0; 5.93 GiB total capacity; 1.23 GiB already allocated; 33.00 MiB free; 1.25 GiB reserved in total by PyTorch)"
Not sure if that's what you want, but you can probably do the following
class MyEfficientNet(nn.Module): def __init__(self): super().__init__() # EfficientNet self.network = EfficientNet.from_pretrained("efficientnet-b0") # Replace last layer self.network._fc = nn.Sequential(nn.Linear(self.network._fc.in_features, 512), nn.ReLU(), nn.Dropout(0.25), nn.Linear(512, 128), nn.ReLU(), nn.Dropout(0.50), nn.Linear(128,classes)) def forward(self, x): out = self.network(x) return out model = MyEfficientNet()And add as much layers you want. as the num of classes, efficientnet has 1000 outputs because of ImageNet, so changing that just make easier for you to train your custom dataset without having to rewrite the last layer.
I tried you code, but starts to get GPU memory error ,even for light b0, before i could use b2 for train
"CUDA out of memory. Tried to allocate 22.00 MiB (GPU 0; 5.93 GiB total capacity; 1.23 GiB already allocated; 33.00 MiB free; 1.25 GiB reserved in total by PyTorch)"
If you are trying to run 224x224 inputs it probably cant fit in 6gb even on mini batches. Try a lower resolution or smaller batch, also check if your forward pass on test is using with torch.no_grad(), one of these could solve your problem, I think its the first case, I can only train inputs close to 64x64 in my 4gb gpu using B0.