Created
June 1, 2019 06:43
-
-
Save InnovArul/53ea0f93462bf1d8f77c9bb0e6e54125 to your computer and use it in GitHub Desktop.
Pytorch inception-v3 query from https://discuss.pytorch.org/t/pytorch-inception-v3/46681
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
from __future__ import print_function | |
from __future__ import division | |
import torch | |
import torch.nn as nn | |
import torch.optim as optim | |
import numpy as np | |
import torchvision | |
from torchvision import datasets, models, transforms | |
from torch.autograd import Variable | |
import matplotlib.pyplot as plt | |
import time | |
import os | |
import copy | |
import torchvision.models as models | |
from tqdm import tqdm | |
inception = models.inception_v3(pretrained=True) | |
for param in inception.parameters(): | |
param.requires_grad = False | |
inception.AuxLogits.fc = nn.Linear(768, 10) | |
inception.fc = nn.Linear(2048, 10) | |
# data augmentation | |
data_transforms = { | |
"train": transforms.Compose( | |
[ | |
transforms.Resize(312), | |
transforms.CenterCrop(299), | |
transforms.ToTensor(), | |
transforms.Lambda(lambda x: x.repeat(3, 1, 1)), | |
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), | |
] | |
), | |
"val": transforms.Compose( | |
[ | |
transforms.Resize(312), | |
transforms.CenterCrop(299), | |
transforms.ToTensor(), | |
transforms.Lambda(lambda x: x.repeat(3, 1, 1)), | |
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), | |
] | |
), | |
} | |
data_dir = "./mnist" | |
num_classes = 10 | |
batch_size = 16 | |
train_loader = torch.utils.data.DataLoader( | |
torchvision.datasets.MNIST( | |
data_dir, train=True, download=True, transform=data_transforms["train"] | |
), | |
batch_size=batch_size, | |
shuffle=True, | |
) | |
test_loader = torch.utils.data.DataLoader( | |
torchvision.datasets.MNIST( | |
data_dir, train=False, download=True, transform=data_transforms["val"] | |
), | |
batch_size=batch_size, | |
shuffle=True, | |
) | |
dataloaders = {"train": train_loader, "val": test_loader} | |
dataset_sizes = {x: len(dataloaders[x].dataset) for x in ["train", "val"]} | |
class_names = train_loader.dataset.classes | |
# train the model | |
# takes the model, dataloaders, criterion, optimizer, device(GPU or CPU) and number of epochs respectively as parameters | |
# returns model, array of validation accuracy, validation loss, train accuracy, train loss respectively | |
def train_model( | |
model, dataloaders, criterion, optimizer, device, num_epochs=25, is_inception=False | |
): | |
since = time.time() | |
val_acc_history = [] | |
best_model_wts = copy.deepcopy(model.state_dict()) | |
best_acc = 0.0 | |
for epoch in range(num_epochs): | |
print("Epoch {}/{}".format(epoch, num_epochs - 1)) | |
print("-" * 10) | |
# Each epoch has a training and validation phase | |
for phase in ["train", "val"]: | |
if phase == "train": | |
model.train() # Set model to training mode | |
else: | |
model.eval() # Set model to evaluate mode | |
running_loss = 0.0 | |
running_corrects = 0 | |
# Iterate over data. | |
print("expecting data") | |
for inputs, labels in tqdm(dataloaders[phase]): | |
inputs = inputs.to(device) | |
labels = labels.to(device) | |
# zero the parameter gradients | |
optimizer.zero_grad() | |
# forward | |
# track history if only in train | |
with torch.set_grad_enabled(phase == "train"): | |
# Get model outputs and calculate loss | |
# Special case for inception because in training it has an auxiliary output. In train | |
# mode we calculate the loss by summing the final output and the auxiliary output | |
# but in testing we only consider the final output. | |
if is_inception and phase == "train": | |
# From https://discuss.pytorch.org/t/how-to-optimize-inception-model-with-auxiliary-classifiers/7958 | |
outputs, aux_outputs = model(inputs) | |
loss1 = criterion(outputs, labels) | |
loss2 = criterion(aux_outputs, labels) | |
loss = loss1 + 0.4 * loss2 | |
else: | |
outputs = model(inputs) | |
loss = criterion(outputs, labels) | |
_, preds = torch.max(outputs, 1) | |
# backward + optimize only if in training phase | |
if phase == "train": | |
loss.backward() | |
optimizer.step() | |
# statistics | |
running_loss += loss.item() * inputs.size(0) | |
running_corrects += torch.sum(preds == labels.data) | |
epoch_loss = running_loss / len(dataloaders[phase].dataset) | |
epoch_acc = running_corrects.double() / len(dataloaders[phase].dataset) | |
print("{} Loss: {:.4f} Acc: {:.4f}".format(phase, epoch_loss, epoch_acc)) | |
# deep copy the model | |
if phase == "val" and epoch_acc > best_acc: | |
best_acc = epoch_acc | |
best_model_wts = copy.deepcopy(model.state_dict()) | |
if phase == "val": | |
val_acc_history.append(epoch_acc) | |
print() | |
time_elapsed = time.time() - since | |
print( | |
"Training complete in {:.0f}m {:.0f}s".format( | |
time_elapsed // 60, time_elapsed % 60 | |
) | |
) | |
print("Best val Acc: {:4f}".format(best_acc)) | |
# load best model weights | |
model.load_state_dict(best_model_wts) | |
return model, val_acc_history | |
# specify loss function | |
criterion = nn.CrossEntropyLoss() | |
# specify optimizer | |
optimizer = torch.optim.Adam(inception.parameters(), lr=0.1) | |
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") | |
inception = inception.to(device) # send the model to the gpu | |
model_name = "inception" | |
model, val_acc = train_model( | |
inception, | |
dataloaders, | |
criterion, | |
optimizer, | |
device, | |
num_epochs=30, | |
is_inception=(model_name == "inception"), | |
) # train model |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment