File size: 5,876 Bytes
727d2ea |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 |
import os
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
from tqdm import tqdm
import matplotlib.pyplot as plt
import timm
# Data augmentation and normalization
transform_train = transforms.Compose([
transforms.RandomResizedCrop(224, scale=(0.8, 1.0)),
transforms.RandomHorizontalFlip(),
transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.2),
transforms.RandomRotation(15),
transforms.RandomAffine(degrees=15, translate=(0.1, 0.1)),
transforms.GaussianBlur(kernel_size=3),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]),
])
transform_val = transforms.Compose([
transforms.Resize(224),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]),
])
# Dataset loading
train_dir = 'D:\\Dataset\\Potato Leaf Disease Dataset in Uncontrolled Environment'
full_ds = datasets.ImageFolder(train_dir, transform=transform_train)
train_size = int(0.8 * len(full_ds))
val_size = len(full_ds) - train_size
train_ds, val_ds = torch.utils.data.random_split(full_ds, [train_size, val_size])
val_ds.dataset.transform = transform_val # Apply validation transforms
train_loader = DataLoader(train_ds, batch_size=32, shuffle=True, num_workers=4)
val_loader = DataLoader(val_ds, batch_size=32, shuffle=False, num_workers=4)
# Device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# Model definition with custom classification head (optional improvement)
model = timm.create_model('mobilenetv3_large_100', pretrained=True)
in_features = model.classifier.in_features
model.classifier = nn.Sequential(
nn.Linear(in_features, 512),
nn.ReLU(),
nn.Dropout(0.3),
nn.Linear(512, len(full_ds.classes))
)
model.to(device)
# Loss and optimizer
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=1e-4)
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='max', factor=0.5, patience=3)
# MixUp augmentation
def mixup_data(x, y, alpha=1.0):
if alpha > 0:
lam = np.random.beta(alpha, alpha)
else:
lam = 1
batch_size = x.size(0)
index = torch.randperm(batch_size).to(x.device)
mixed_x = lam * x + (1 - lam) * x[index, :]
y_a, y_b = y, y[index]
return mixed_x, y_a, y_b, lam
def mixup_criterion(criterion, pred, y_a, y_b, lam):
return lam * criterion(pred, y_a) + (1 - lam) * criterion(pred, y_b)
# Training function with MixUp
def train_epoch(model, train_loader, criterion, optimizer):
model.train()
running_loss, correct_preds, total_preds = 0.0, 0, 0
for inputs, labels in tqdm(train_loader, desc="Training Epoch", leave=False):
inputs, labels = inputs.to(device), labels.to(device)
inputs, targets_a, targets_b, lam = mixup_data(inputs, labels, alpha=1.0)
optimizer.zero_grad()
outputs = model(inputs)
loss = mixup_criterion(criterion, outputs, targets_a, targets_b, lam)
loss.backward()
optimizer.step()
_, preds = torch.max(outputs, 1)
correct_preds += (lam * preds.eq(targets_a).sum().item()
+ (1 - lam) * preds.eq(targets_b).sum().item())
total_preds += labels.size(0)
running_loss += loss.item()
return running_loss / len(train_loader), correct_preds / total_preds
# Validation function
def validate_epoch(model, val_loader, criterion):
model.eval()
running_loss, correct_preds, total_preds = 0.0, 0, 0
with torch.no_grad():
for inputs, labels in tqdm(val_loader, desc="Validating Epoch", leave=False):
inputs, labels = inputs.to(device), labels.to(device)
outputs = model(inputs)
loss = criterion(outputs, labels)
_, preds = torch.max(outputs, 1)
correct_preds += (preds == labels).sum().item()
total_preds += labels.size(0)
running_loss += loss.item()
return running_loss / len(val_loader), correct_preds / total_preds
# Plotting
def plot_metrics(train_loss, val_loss, train_acc, val_acc):
epochs = range(1, len(train_loss) + 1)
plt.figure(figsize=(12, 5))
plt.subplot(1, 2, 1)
plt.plot(epochs, train_loss, label='Training Loss')
plt.plot(epochs, val_loss, label='Validation Loss')
plt.xlabel('Epochs')
plt.ylabel('Loss')
plt.legend()
plt.subplot(1, 2, 2)
plt.plot(epochs, train_acc, label='Training Accuracy')
plt.plot(epochs, val_acc, label='Validation Accuracy')
plt.xlabel('Epochs')
plt.ylabel('Accuracy')
plt.legend()
plt.show()
# Training loop
num_epochs = 20
train_losses, val_losses = [], []
train_accuracies, val_accuracies = [], []
for epoch in range(num_epochs):
print(f"\nEpoch {epoch+1}/{num_epochs}")
train_loss, train_acc = train_epoch(model, train_loader, criterion, optimizer)
val_loss, val_acc = validate_epoch(model, val_loader, criterion)
scheduler.step(val_acc)
print(f"Train Loss: {train_loss:.4f}, Accuracy: {train_acc:.4f}")
print(f"Val Loss: {val_loss:.4f}, Accuracy: {val_acc:.4f}")
train_losses.append(train_loss)
val_losses.append(val_loss)
train_accuracies.append(train_acc)
val_accuracies.append(val_acc)
plot_metrics(train_losses, val_losses, train_accuracies, val_accuracies)
best_val_acc = 0.0
save_path = 'D:\\Dataset\\Potato Leaf Disease Dataset in Uncontrolled Environment\\best_model.pth'
os.makedirs(os.path.dirname(save_path), exist_ok=True)
if val_acc > best_val_acc:
best_val_acc = val_acc
torch.save(model.state_dict(), save_path)
print(f"✅ Best model saved with val_acc: {val_acc:.4f}")
|