From d37d811f66a447d3ad1aa91fc9a28e42cdd76aac Mon Sep 17 00:00:00 2001 From: krataratha Date: Sat, 20 Jun 2026 20:23:44 +0530 Subject: [PATCH] Enhance Detector class with feature extraction and accuracy Added feature to extract raw embedding features and calculate batch accuracy during training. --- src/model.py | 44 +++++++++++++++++++++++--------------------- 1 file changed, 23 insertions(+), 21 deletions(-) diff --git a/src/model.py b/src/model.py index 00e138f..ee6b7bc 100644 --- a/src/model.py +++ b/src/model.py @@ -5,35 +5,37 @@ from efficientnet_pytorch import EfficientNet from utils.sam import SAM - - class Detector(nn.Module): def __init__(self): super(Detector, self).__init__() - self.net=EfficientNet.from_pretrained("efficientnet-b4",advprop=True,num_classes=2) - self.cel=nn.CrossEntropyLoss() - self.optimizer=SAM(self.parameters(),torch.optim.SGD,lr=0.001, momentum=0.9) - - + self.net = EfficientNet.from_pretrained("efficientnet-b4", advprop=True, num_classes=2) + self.cel = nn.CrossEntropyLoss() + self.optimizer = SAM(self.parameters(), torch.optim.SGD, lr=0.001, momentum=0.9) - def forward(self,x): - x=self.net(x) - return x + def forward(self, x, return_features=False): + # Feature: Added toggle to extract raw embedding features before the classifier head + if return_features: + return self.net.extract_features(x) + return self.net(x) - def training_step(self,x,target): + def training_step(self, x, target): for i in range(2): - pred_cls=self(x) - if i==0: - pred_first=pred_cls - loss_cls=self.cel(pred_cls,target) - loss=loss_cls + pred_cls = self(x) + if i == 0: + pred_first = pred_cls + + loss = self.cel(pred_cls, target) self.optimizer.zero_grad() loss.backward() - if i==0: - self.optimizer.first_step(zero_grad=True) - else: - self.optimizer.second_step(zero_grad=True) + + # Dynamically select the SAM step based on the loop index + step_fn = self.optimizer.first_step if i == 0 else self.optimizer.second_step + step_fn(zero_grad=True) + + # Feature: Now calculates and returns batch accuracy alongside the predictions + acc = (pred_first.argmax(dim=-1) == target).float().mean() + return pred_first, loss, acc return pred_first - \ No newline at end of file +