EfficientNet-PyTorch
EfficientNet-PyTorch copied to clipboard
Get the number of output features from the last layer and remove the last layer?
Hello,
I am trying to add the LoRA layer to the EfficientNet-B0 till B7 algorithms. However, I am not succeed in getting the number of output features and remove the last fully connected layer. Can you help me?
I tried the following code, but there is no "._fc" attribute on Python pre-trained algorithm.
Define EfficientNet-B0 with LoRA
class EfficientNetB0LoRA(nn.Module):
def __init__(self, num_classes, lora_rank):
super(EfficientNetB0LoRA, self).__init__()
# Load pre-trained EfficientNet-B0 model
self.efficientnet_b0 = EfficientNet.from_pretrained('efficientnet-b0')
# Get number of input features for the LoRALayer
num_features = self.efficientnet_b0._fc.in_features
# Replace the classifier with an identity layer
self.efficientnet_b0._fc = nn.Identity()
# Add LoRA layer
self.lora = LoRALayer(num_features, num_classes, lora_rank)
def forward(self, x):
x = self.efficientnet_b0(x)
x = self.lora(x)
return x