Homec4science
Diffusion LPBF Transfer Learning (master) Browse

R11778/Transfer_Learning_VGG16.py
0a854839b33emaster

Last Change
Raw File

Transfer_Learning_VGG16.py
View Options

# -*- coding: utf-8 -*-
"""
Created on Sat Feb 8 22:10:18 2020
@author: srpv
"""
import torchvision.transforms as transforms
import torchvision
import torch
from torchsummary import summary
import matplotlib.pyplot as plt
import numpy as np
from torch.optim.lr_scheduler import StepLR
from mlxtend.plotting import plot_confusion_matrix
import seaborn as sns
from torchvision import datasets
from Heatmap import heatmap , annotate_heatmap
#torch.cuda.empty_cache()
from torch import optim, cuda
import os
from PIL import Image
import pandas as pd
import torchvision.models as models
from torch import nn
from collections import OrderedDict
# Whether to train on a gpu
train_on_gpu = cuda.is_available()
print(f'Train on gpu: {train_on_gpu}')
from Utils import *
#%%
classes = ('Balling', 'LoF', 'Nopores','Keyhole')
PATH = './VGG16-Pytorch.pth'
Trained_model = torch.load(PATH)
print(Trained_model)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print(device)
datadir = 'Bronze_dataset/'
traindir = datadir + 'Train/'
testdir = datadir + 'Test/'
#%%
for name, child in Trained_model.named_children():
if name == 'classifier':
for name2, params in child.named_parameters():
print(name,name2)
params.requires_grad = True
else:
for name2, params in child.named_parameters():
params.requires_grad = False
Trained_model.to(device)
summary(Trained_model, (3 ,32 ,32))
#%%
def get_lr(optimizer):
for param_group in optimizer.param_groups:
print('Learning rate =')
print(param_group['lr'])
return param_group['lr']
transform = transforms.Compose([transforms.Resize((512,512)),
transforms.ToTensor()])
#transform = transforms.Compose([transforms.ToTensor()])
trainload = datasets.ImageFolder(root=traindir, transform=transform)
trainset = torch.utils.data.DataLoader(trainload, batch_size=10,
shuffle=True, num_workers=0)
testload = datasets.ImageFolder(root=testdir, transform=transform)
testset = torch.utils.data.DataLoader(testload, batch_size=10,
shuffle=True, num_workers=0)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
#device = torch.device("cpu")
print(device)
net= Trained_model
net.to(device)
costFunc = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(net.parameters(),lr=0.001,momentum=0.9)
scheduler = StepLR(optimizer, step_size = 25, gamma= 0.5 )
Loss_value =[]
Iteration_count=1
iteration=[]
Epoch_count=0
Total_Epoch =[]
Accuracy=[]
Learning_rate=[]
for epoch in range(2):
learingrate_value = get_lr(optimizer)
Learning_rate.append(learingrate_value)
closs = 0
scheduler.step()
for i,batch in enumerate(trainset,0):
data,output = batch
data,output = data.to(device),output.to(device)
prediction = net(data)
loss = costFunc(prediction,output)
# print("loss",loss)
closs += loss
# print("closs",closs)
optimizer.zero_grad()
loss.backward()
optimizer.step()
Iteration_count = Iteration_count + i
iteration.append(Iteration_count)
# closs = loss.item()
#print every 1000th time
if i%100 == 0:
print('[%d %d] loss: %.4f'% (epoch+1,i+1,loss))
loss_train = closs / i
print('Loss on epoch: [%d] is %.4f'% (epoch+1,loss_train))
Loss_value.append(loss_train)
correctHits=0
total=0
for batches in testset:
data,output = batches
data,output = data.to(device),output.to(device)
prediction = net(data)
# _,prediction = torch.max(prediction.data,1) #returns max as well as its index
prediction = torch.argmax(prediction, dim=1)
total += output.size(0)
correctHits += (prediction==output).sum().item()
Epoch_accuracy = (correctHits/total)*100
Accuracy.append(Epoch_accuracy)
print('Accuracy on epoch ',epoch+1,'= ',str((correctHits/total)*100))
Epoch_count = epoch+1
Total_Epoch.append (Epoch_count)
y_pred = []
y_true = []
correctHits=0
total=0
for batches in testset:
data,output = batches
data,output = data.to(device),output.to(device)
prediction = net(data)
# _,prediction = torch.max(prediction.data,1) #returns max as well as its index
prediction = torch.argmax(prediction, dim=1)
total += output.size(0)
correctHits += (prediction==output).sum().item()
prediction=prediction.data.cpu().numpy()
output=output.data.cpu().numpy()
y_true.extend(output) # Save Truth
y_pred.extend(prediction)
print('Accuracy = '+str((correctHits/total)*100))
print('Finished Training')
PATH = './VGG16Bronze-Pytorch.pth'
torch.save(net.state_dict(), PATH)
torch.save(net, PATH)
#Trained_model = torch.load(PATH)
#%%
Loss_value= torch.stack(Loss_value)
Loss_value=Loss_value.cpu().detach().numpy()
plots(iteration,Loss_value,Total_Epoch,Accuracy,Learning_rate,'VGG16_Transfer_Bronze')
count_parameters(net)
plotname= 'VGG16'+'_Transfer_Bronze'+'_confusion_matrix'+'.png'
plot_confusion_matrix(y_true, y_pred,classes,plotname)
c4science · Help