Overview

This project uses EfficientNet-B0 for terrain classification, distinguishing between different terrain types such as mountains, forests, glaciers, deserts, and coasts. Large datasets of terrain data, collected from photographs and satellite images, need rapid labeling for efficient processing. With the use of Transfer Learning and GPU acceleration, the model offers fast inference with high accuracy, making it suitable for handling large-scale datasets effectively.

Key features include:

• Data Augmentation for Robust Training
• GPU-accelerated Training for Speed
• Model Evaluation using Confusion Matrices
• Early Stopping to Prevent Overfitting

Code Walkthrough

1. Importing Necessary Modules

We start by importing all the required modules, including PyTorch for deep learning, EfficientNet from torchvision, and data loaders for handling datasets.

# Importing required libraries
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, models, transforms
from torch.utils.data import DataLoader
from torchvision.models import EfficientNet_B0_Weights
import matplotlib.pyplot as plt
    

2. Checking GPU Availability

This step ensures that the model utilizes GPU if available for faster training and inference.

# Check if a GPU is available
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")
    

3. Data Augmentation and Loading

We define data transformations to improve the model's robustness and load the datasets using PyTorch's ImageFolder.

# Define data augmentation techniques
data_transforms = transforms.Compose([
    transforms.RandomResizedCrop(224),
    transforms.RandomHorizontalFlip(),
    transforms.ColorJitter(brightness=0.2, contrast=0.2),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
])

# Load training and testing datasets
train_dataset = datasets.ImageFolder(root="train_data_path", transform=data_transforms)
test_dataset = datasets.ImageFolder(root="test_data_path", transform=data_transforms)

# Create data loaders
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)
    

4. Modifying the EfficientNet-B0 Model

We modify the classifier layer to match the number of terrain classes (5 classes in this case).

# Load EfficientNet-B0 with pretrained weights
model = models.efficientnet_b0(weights=EfficientNet_B0_Weights.DEFAULT)

# Modify the final classifier layer
num_classes = 5
model.classifier[1] = nn.Linear(model.classifier[1].in_features, num_classes)

# Move the model to the appropriate device
model = model.to(device)
    

5. Defining Loss Function and Optimizer

We use CrossEntropyLoss as the loss function and Adam as the optimizer for efficient learning.

# Define loss function and optimizer
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.0001)
    

6. Implementing Early Stopping and Learning Rate Scheduler

A learning rate scheduler adjusts the learning rate dynamically, and early stopping prevents overfitting.

# Learning rate scheduler
scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, factor=0.5, patience=2, verbose=True)
    

7. Training the Model

We define the training loop, track accuracy, and implement early stopping.

def train_model(model, train_loader, criterion, optimizer, scheduler, num_epochs=15, patience=5):
    best_acc = 0.0
    patience_counter = 0

    for epoch in range(num_epochs):
        model.train()
        running_loss = 0.0
        correct, total = 0, 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()

            running_loss += loss.item()
            _, preds = torch.max(outputs, 1)
            correct += (preds == labels).sum().item()
            total += labels.size(0)

        epoch_acc = correct / total * 100
        print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {running_loss:.4f}, Accuracy: {epoch_acc:.2f}%')

        scheduler.step(epoch_acc)

        if epoch_acc > best_acc:
            best_acc = epoch_acc
            patience_counter = 0
            torch.save(model.state_dict(), "best_model.pth")
        else:
            patience_counter += 1
            if patience_counter >= patience:
                print("Early stopping triggered.")
                break
    

8. Evaluating the Model

Finally, we evaluate the trained model on the test dataset to check its performance.

def evaluate_model(model, test_loader):
    model.eval()
    correct, total = 0, 0
    with torch.no_grad():
        for images, labels in test_loader:
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            _, preds = torch.max(outputs, 1)
            correct += (preds == labels).sum().item()
            total += labels.size(0)

    print(f'Test Accuracy: {100 * correct / total:.2f}%')

# Evaluate the trained model
evaluate_model(model, test_loader)