EfficientNet_b0 with Lightning

Rami_Abdelouahab: <p>i m getting low probs on predicted pics compares to the model without lightining.<br>
Here is my code:</p>
<pre><code class="lang-auto">class EfficientNet(L.LightningModule):
    def __init__(self,num_classes):
        super().__init__()
        self.weights = torchvision.models.EfficientNet_B0_Weights.DEFAULT
        
#         self.model = torchvision.models.efficientnet_b0(weights=self.weights)
        self.model =  torchvision.models.efficientnet_b0(pretrained=True)
        
        self.loss = nn.CrossEntropyLoss()
        
        self.accuracy = Accuracy(task="multiclass", num_classes=num_classes)
        
        
        for param in self.model.features.parameters():
            param.requires_grad = False
            
        self.model.classifier = nn.Sequential(
                nn.Dropout(p=0.2,inplace=True),
                nn.Linear(in_features=1280,out_features=num_classes)
                
        )
        

        
        
    def forward(self,x):
        x = self.model(x)
        return x
    
    def configure_optimizers(self):
        
        params = self.model.parameters()
        optimizer = torch.optim.Adam(params=params,lr=0.001)
        
        return optimizer
    
    def training_step(self, batch, batch_idx):
        # training_step defines the train loop.
        x, y = batch 
        logits = self.forward(x)
        
        loss = self.loss(logits,y)
    
        
        acc = self.accuracy(logits, y)
        
        self.log('train_loss', loss, prog_bar=True, logger=True)
        self.log('train_acc', acc, prog_bar=True, logger=True)

    
        return {'loss':loss,'acc':acc}
    
#     def validation_step(self, batch, batch_idx):
#         results = self.training_step(batch,batch_idx)
#         return results
    
    def on_train_epoch_end(self):
        
      
        self.log('train_acc_epoch', self.accuracy)
        self.log("step", self.current_epoch)
    
#     def on_validation_epoch_end(self):
        
# #         avg_val_loss = torch.tensor([x['loss'] for x in self.val_step_outputs]).mean()
# #         avg_val_acc = torch.tensor([x['acc'] for x in self.val_step_outputs]).mean()
        
#         pbar = {'avg_val_acc':self.accuracy}
#         # log epoch metric
#         self.log('avg_val_acc', self.accuracy)
        
#         return {'progress_bar':pbar}
    
    def train_dataloader(self):
        
            train_data = datasets.ImageFolder(TRAIN_DIR, transform=manual_transform)
            
            NUM_WORKERS = os.cpu_count()
            
            # Turn images into data loaders
            train_dataloader = DataLoader(
              train_data,
              batch_size=32,
              shuffle=True,
              num_workers=NUM_WORKERS,
              pin_memory=True)
    
            return train_dataloader
    
    def val_dataloader(self):
        
            test_data = datasets.ImageFolder(TEST_DIR, transform=manual_transform)
            
            NUM_WORKERS = os.cpu_count()
            
            # Turn images into data loaders
            test_dataloader = DataLoader(
              test_data,
              batch_size=32,
              shuffle=False,
              num_workers=NUM_WORKERS,
              pin_memory=True)
    
            return test_dataloader
        
        
model = EfficientNet(len(class_names))
</code></pre>