roi-triain: Use resnet50

research/roi-train.py

@@ -10,6 +10,8 @@ from datetime import datetime
 from collections import defaultdict
 import random
 import argparse
+from kornia.losses.focal import FocalLoss
+from torchvision.models import ResNet50_Weights
 
 class SideBySideDataset(Dataset):
     def __init__(self, img_dir, labels_file, transform=None):
@@ -81,8 +83,7 @@ def main():
     train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
     val_loader = DataLoader(val_dataset, batch_size=32, shuffle=True)
 
-    # 加载预训练的ResNet18模型
-    model = models.resnet18(pretrained=True)
+    model = models.resnet50(weights=ResNet50_Weights.IMAGENET1K_V1)
 
     # 修改最后一层全连接层，使其输出为2（二分类）
     num_ftrs = model.fc.in_features
@@ -92,7 +93,9 @@ def main():
     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
     model = model.to(device)
 
-    criterion = nn.CrossEntropyLoss()
+    # criterion = nn.CrossEntropyLoss()
+    pos_weight = 0.9
+    criterion = FocalLoss(0.25, weight=torch.Tensor([1 - pos_weight, pos_weight]))
     optimizer = optim.Adam(model.parameters(), lr=0.0001)
     #optimizer = torch.optim.SGD( model.parameters(), lr=0.0001, momentum=0.09, weight_decay=1e-4)
 
@@ -109,13 +112,14 @@ def main():
         model.train()
         running_loss = 0.0
         for images, labels in train_loader:
-            images, labels = images.to(device), labels.to(device)
+            # images, labels = images.to(device), labels.to(device)
             optimizer.zero_grad()
             outputs = model(images)
             loss = criterion(outputs, labels)
-            loss.backward()
+            loss_sum = loss.sum()
+            loss_sum.backward()
             optimizer.step()
-            running_loss += loss.item()
+            running_loss += loss_sum
 
         print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {running_loss/len(train_loader):.4f}')
 
@@ -123,29 +127,34 @@ def main():
 
         # 验证阶段
         model.eval()
-        correct = 0
-        total = 0
+        pos_correct = 0
+        pos_total = 0
+        neg_correct = 0
+        neg_total = 0
         with torch.no_grad():
             for images, labels in val_loader:
                 images, labels = images.to(device), labels.to(device)
                 outputs = model(images)
                 _, predicted = torch.max(outputs.data, 1)
-                total += labels.size(0)
-                correct += (predicted == labels).sum().item()
+                pos_correct += ((predicted == labels) & (labels == 1)).sum().item()
+                pos_total += (labels == 1).sum().item()
+                neg_correct += ((predicted == labels) & (labels == 0)).sum().item()
+                neg_total += (labels == 0).sum().item()
 
-        last_accu = 100 * correct / total
+        last_accu = 100 * (pos_correct + neg_correct) / (pos_total + neg_total)
         prev_accu.append(last_accu)
-        if epoch > 5:
+        if 0 and epoch > 5:
             avg = sum(prev_accu[-5:]) / 5
             variance = sum((x - avg) ** 2 for x in prev_accu[-5:]) / 5
             print(f"variance={variance:.4f}")
             if variance < 1 and not (prev_accu[-1] > prev_accu[-2] and prev_accu[-2] > prev_accu[-3] and prev_accu[-3] > prev_accu[-4]):
                 print(f"Early stopping condition met: {variance:.4f}")
                 break
-        print(f'Validation Accuracy: {last_accu:.2f}%')
+        print(f'Pos Accu: {100 * pos_correct / pos_total:.2f}%')
+        print(f'Neg Accu: {100 * neg_correct / neg_total:.2f}%')
 
     dt = datetime.now().strftime("%Y%m%d_%H%M%S")
-    torch.save(model.state_dict(), f'data/roi/models/resnet18_{dt}_{last_accu:.2f}.pth')
+    torch.save(model.state_dict(), f'data/roi/models/resnet_{dt}_{last_accu:.2f}.pth')
 
 if __name__ == "__main__":
     main()