alexnet.py

# -*- coding: utf-8 -*-
"""AlexNet.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1ABqe0qMluk3sfqmwdxABpyUl93efb4Al
"""

from google.colab import drive
drive.mount('/content/drive')

pip install cnn_finetune

import skimage
import pandas
from __future__ import print_function, division
import os
import torch
import pandas as pd
from skimage import io, transform
import numpy as np
import matplotlib.pyplot as plt
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms, utils
from PIL import Image
import pandas as pd
# Ignore warnings
import warnings
warnings.filterwarnings("ignore")
from torch.utils.data.sampler import SubsetRandomSampler
from torch.utils.data import DataLoader

import torch.nn as nn
import torch.nn.functional as F
import argparse
import torch
import torchvision
import torchvision.transforms as transforms
from torch.autograd import Variable
import torch.nn as nn
import torch.optim as optim

from cnn_finetune import make_model

parser = argparse.ArgumentParser(description='cnn_finetune')
parser.add_argument('-f')
parser.add_argument('--batch-size', type=int, default=16, metavar='N',
                    help='input batch size for training (default: 4)')
parser.add_argument('--epochs', type=int, default=30, metavar='N',
                    help='number of epochs to train (default: 30)')
parser.add_argument('--lr', type=float, default=0.001, metavar='LR',
                    help='learning rate (default: 0.001)')
parser.add_argument('--momentum', type=float, default=0.9, metavar='M',
                    help='SGD momentum (default: 0.9)')
parser.add_argument('--no-cuda', action='store_true', default=False,
                    help='disables CUDA training')
parser.add_argument('--model-name', type=str, default='alexnet', metavar='M',
                    help='model name (default: alexnet)')

args = parser.parse_args()
use_cuda = not args.no_cuda and torch.cuda.is_available()
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

class FloodTinyDataset(Dataset):

    def __init__(self, df, transform = None):
       
        self.flood_tiny_data = df
        ''' Change self.root_dir to load images '''
        # self.root_dir = f"/content/drive/My Drive/ladi/Images/flood_tiny/"
        self.transform = transform

    def __len__(self):
        return len(self.flood_tiny_data)

    def __getitem__(self, idx):
        if torch.is_tensor(idx):
            idx = idx.tolist()
            
        # pos = self.flood_tiny_data.iloc[idx, 9].rfind('/')+1
        # img_name = os.path.join(self.root_dir, self.flood_tiny_data.iloc[idx, 9][pos:])  
        img_name = self.flood_tiny_data.iloc[idx, 10]

        image = Image.fromarray(io.imread(img_name))
        uuid = self.flood_tiny_data.iloc[idx, 1]
        timestamp = self.flood_tiny_data.iloc[idx, 2]
        gps_lat = self.flood_tiny_data.iloc[idx, 3]
        gps_lon = self.flood_tiny_data.iloc[idx, 4]
        gps_alt = self.flood_tiny_data.iloc[idx, 5]
        file_size = self.flood_tiny_data.iloc[idx, 6]
        width = self.flood_tiny_data.iloc[idx, 7]
        height = self.flood_tiny_data.iloc[idx, 8]
        ### Labels should be numerical, not bool for training ###
        if self.flood_tiny_data.iloc[idx, -1] == True:
            label = 1
        else:
            label = 0
        
        if self.transform:
            image = self.transform(image)

        sample = {'image': image, 'image_name': img_name, 'label': label, 'uuid': uuid, 'timestamp': timestamp, 'gps_lat': gps_lat, 'gps_lon': gps_lon, 'gps_alt': gps_alt, 'orig_file_size': file_size, 'orig_width': width, 'orig_height': height}

        return sample

csv_file = '/content/drive/My Drive/DS440_Model/flood_tiny_metadata_L.csv'
label_csv = '/content/drive/My Drive/DS440_Model/flood_tiny_label_L.csv'

metadata = pd.read_csv(csv_file)
label = pd.read_csv(label_csv)
all_data = pd.merge(metadata, label, on="s3_path")

class_zero = all_data[all_data['label']==False]
class_one = all_data[all_data['label']==True]
flood_tiny_data = pd.concat([class_zero, class_one]).sample(frac=1).reset_index(drop=True)

flood_tiny_dataset = FloodTinyDataset(df=flood_tiny_data)

transformed_dataset = FloodTinyDataset(df=flood_tiny_data, transform=transforms.Compose([transforms.Resize(256),
transforms.RandomRotation(10),
transforms.RandomCrop(256),
transforms.RandomHorizontalFlip(), 
transforms.ToTensor()]))

batch_size=args.batch_size
test_split_ratio = .2
shuffle_dataset = True
random_seed= 42
# Creating data indices for training and validation splits:
dataset_size = len(transformed_dataset)
indices = list(range(dataset_size))
split = int(np.floor(test_split_ratio * dataset_size))
if shuffle_dataset :
    np.random.seed(random_seed)
    np.random.shuffle(indices)
train_indices, test_indices = indices[split:], indices[:split]

train_sampler = SubsetRandomSampler(train_indices)
test_sampler = SubsetRandomSampler(test_indices)

train_loader = torch.utils.data.DataLoader(transformed_dataset, batch_size=batch_size,
                                           sampler=train_sampler)
test_loader = torch.utils.data.DataLoader(transformed_dataset, batch_size=batch_size,
                                                sampler=test_sampler)

def train(model, epoch, optimizer, train_loader, criterion=nn.CrossEntropyLoss()):
    running_loss = 0
    total_size = 0
    model.train()
    for i, data in enumerate(train_loader, 0):

        # get the inputs; data is a list of [inputs, labels]
        inputs = data['image']
        labels = data['label']
        inputs = inputs.to(device)
        labels = labels.to(device)
        # casting int to long for loss calculation#
        labels = labels.long()

        # zero the parameter gradients
        optimizer.zero_grad()

        # forward + backward + optimize
        outputs = model(inputs)
        loss = criterion(outputs, labels)

        running_loss += loss.item()
        total_size += inputs.size(0)
        loss.backward()
        optimizer.step()

        if i % 20 == 19:   
            print('[%d, %3d] loss: %.3f' %
                  (epoch + 1, i + 1, running_loss / 20))
            running_loss = 0.0

print('Finished Training')

def test(model, test_loader, criterion=nn.CrossEntropyLoss()):
    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for data in test_loader:
            inputs = data['image']
            labels = data['label']
            inputs = inputs.to(device)
            labels = labels.to(device)
            
            outputs = model(inputs)
            
            _, predicted = torch.max(outputs.data, 1)
            #test_loss += criterion(output, target).item()
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

            #correct += pred.eq(target.data.view_as(pred)).long().cpu().sum().item()
    accuracy = 100 * (correct / total)
    print('Accuracy of the network on test images: %d %%' % (accuracy))
    state = {'epoch': epoch, 'state_dict': model.state_dict(),
             'optimizer_state_dict': optimizer.state_dict()}
    model_name = 'Alexnet_%d_%d.pth' % (epoch, accuracy)
    PATH = f"/content/drive/My Drive/{model_name}"
    torch.save(state, PATH)

model_name = args.model_name

# classes = ('0','1') 
model = make_model(
        model_name,
        pretrained=True,
        num_classes=2,
        input_size= (256,256),
    )
model = model.to(device)
optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum= args.momentum)
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.1)
for epoch in range(0, args.epochs):
  scheduler.step(epoch)
  train(model, epoch, optimizer, train_loader)
  test(model, test_loader)