i m getting low probs on predicted pics compares to the model without lightining.
Here is my code:
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))