main.py

import copy
import random
import numpy as np
import os
from tqdm import tqdm
import torch
import torch.backends.cudnn as cudnn
import torch.nn as nn
import torch.optim
from torch.utils.data import DataLoader
import torchvision.models as models
from helper import TqdmUpTo, MetricMonitor, calculate_accuracy
from dataset.utils import get_train_test
from dataset.dataset import SmokeDataset
from visualize import display_image_grid, visualize_augmentations, plot_curves
from config import Config
from collections import defaultdict

import warnings
warnings.filterwarnings('ignore')


cudnn.benchmark = True


def train(train_loader, model, criterion, optimizer, epoch, cfg):
    metric_monitor = MetricMonitor()
    model.train()
    stream = tqdm(train_loader, position=0, leave=True, colour='green')
    for i, (images, target) in enumerate(stream, start=1):
        images = images.to(cfg.device, non_blocking=True)
        target = target.to(
            cfg.device, non_blocking=True).float().view(-1, 1)
        output = model(images)
        loss = criterion(output, target)
        accuracy = calculate_accuracy(output, target)
        metric_monitor.update("Loss", loss.item())
        metric_monitor.update("Accuracy", accuracy)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        stream.set_description(
            "Epoch: {epoch}. Train. {metric_monitor}".format(
                epoch=epoch, metric_monitor=metric_monitor)
        )
    return metric_monitor

def validate(val_loader, model, criterion, epoch, cfg):
    metric_monitor = MetricMonitor()
    model.eval()
    stream = tqdm(val_loader, position=0, leave=True, colour='red')
    with torch.no_grad():
        for i, (images, target) in enumerate(stream, start=1):
            images = images.to(cfg.device, non_blocking=True)
            target = target.to(cfg.device, non_blocking=True).float().view(-1, 1)
            output = model(images)
            loss = criterion(output, target)
            accuracy = calculate_accuracy(output, target)

            metric_monitor.update("Loss", loss.item())
            metric_monitor.update("Accuracy", accuracy)
            stream.set_description(
                "Epoch: {epoch}. Validation. {metric_monitor}".format(epoch=epoch, metric_monitor=metric_monitor)
            )
    return metric_monitor

def main():
    cfg = Config()
    
    history = defaultdict()
    history['val_loss'], history['train_loss'] = list(), list()
    history['val_acc'], history['train_acc'] = list(), list()
    
    
    print('Creating training, validation set')
    (train_images, train_labels), (val_images, val_labels), (test_images,
                                                             test_labels) = get_train_test(cfg.dataset_dir, validation_size=0.1)
    #print(len(train_images), len(val_images), len(test_images))

    ## Visualize the training images
    ## some random samples are shown with the labels
    #display_image_grid(train_images, train_labels)

    train_dataset = SmokeDataset(
        train_images, train_labels, resize=cfg.resize, transform=cfg.train_transform)

    val_dataset = SmokeDataset(
        val_images, val_labels, resize=cfg.resize, transform=cfg.val_transform)

    ## Visualize the augmented images after augmentation
    #visualize_augmentations(train_dataset, idx=220, name='augments.jpg')

    train_loader = DataLoader(
        train_dataset, batch_size=cfg.batch_size, shuffle=True, num_workers=cfg.num_workers, pin_memory=cfg.pin_memory,
    )


    val_loader = DataLoader(
        val_dataset, batch_size=cfg.batch_size, shuffle=True, num_workers=cfg.num_workers, pin_memory=cfg.pin_memory,
    )

    model = getattr(models, cfg.model_name)(
        pretrained=False, num_classes=cfg.num_classes,)
    
    if os.path.isfile(cfg.model_path):
        model.load_state_dict(torch.load(cfg.model_path, map_location=cfg.device))
        print('Saved Model found at',cfg.model_path ,' and loaded')
            
    model = model.to(cfg.device)
    criterion = nn.BCEWithLogitsLoss().to(cfg.device)
    optimizer = torch.optim.Adam(model.parameters(), lr=cfg.learning_rate)

    BEST_LOSS = np.inf
    
    for epoch in range(1, cfg.epochs + 1):
        hist = train(train_loader, model, criterion, optimizer, epoch, cfg)
        history['train_loss'].append(hist.get('Loss'))
        history['train_acc'].append(hist.get('Accuracy'))
        
        
        hist = validate(val_loader, model, criterion, epoch, cfg)
        if hist.get('Loss') < BEST_LOSS:
            print('Saving model...')
            torch.save(model.state_dict(), cfg.model_path)
            BEST_LOSS = hist.get('Loss')
        history['val_loss'].append(hist.get('Loss'))
        history['val_acc'].append(hist.get('Accuracy'))
        
    plot_curves(history['train_acc'], history['val_acc'], ylabel='Accuracy', name='accuracy.png')
    plot_curves(history['train_loss'], history['val_loss'], ylabel='Loss', name='loss.png')
    
    
        
    
    
    
        
    
    
if __name__ == '__main__':
    main()