File size: 2,993 Bytes
a3fdab1 |
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 |
import os
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, random_split
from torchvision.datasets import ImageFolder
from torchvision import transforms
from models.cnn import CNNModel
from utils.transforms import get_transforms
def train_model(data_dir='data/intel/seg_train', epochs=10, batch_size=32, save_path='saved_models/cnn_model.pth'):
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# Load dataset with ImageFolder (expects folder structure: class subfolders)
full_dataset = ImageFolder(root=data_dir, transform=get_transforms(train=True))
class_names = full_dataset.classes
print(f"Classes: {class_names}")
# Split dataset 80% train, 20% val
train_size = int(0.8 * len(full_dataset))
val_size = len(full_dataset) - train_size
train_ds, val_ds = random_split(full_dataset, [train_size, val_size])
# Update val transforms (no augmentation)
val_ds.dataset.transform = get_transforms(train=False)
train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True)
val_loader = DataLoader(val_ds, batch_size=batch_size)
model = CNNModel(num_classes=len(class_names)).to(device)
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
for epoch in range(epochs):
model.train()
total_loss = 0
total_correct = 0
for images, labels in train_loader:
images, labels = images.to(device), labels.to(device)
optimizer.zero_grad()
outputs = model(images)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
total_loss += loss.item() * images.size(0)
total_correct += (outputs.argmax(1) == labels).sum().item()
train_loss = total_loss / len(train_loader.dataset)
train_acc = total_correct / len(train_loader.dataset)
# Validation
model.eval()
val_loss = 0
val_correct = 0
with torch.no_grad():
for images, labels in val_loader:
images, labels = images.to(device), labels.to(device)
outputs = model(images)
loss = criterion(outputs, labels)
val_loss += loss.item() * images.size(0)
val_correct += (outputs.argmax(1) == labels).sum().item()
val_loss /= len(val_loader.dataset)
val_acc = val_correct / len(val_loader.dataset)
print(f"Epoch {epoch+1}/{epochs} — Train loss: {train_loss:.4f}, Train acc: {train_acc:.4f}, Val loss: {val_loss:.4f}, Val acc: {val_acc:.4f}")
os.makedirs(os.path.dirname(save_path), exist_ok=True)
torch.save({
'model_state_dict': model.state_dict(),
'class_names': class_names
}, save_path)
print(f"Model saved to {save_path}")
if __name__ == "__main__":
train_model()
|