Page Menu
Home
c4science
Search
Configure Global Search
Log In
Files
F83469065
Main ResNet18.py
No One
Temporary
Actions
Download File
Edit File
Delete File
View Transforms
Subscribe
Mute Notifications
Award Token
Subscribers
None
File Metadata
Details
File Info
Storage
Attached
Created
Tue, Sep 17, 07:39
Size
5 KB
Mime Type
text/x-python
Expires
Thu, Sep 19, 07:39 (2 d)
Engine
blob
Format
Raw Data
Handle
20759242
Attached To
R11778 LPBF Transfer Learning
Main ResNet18.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'
)
#%%
datadir
=
'Stainless_steel_dataset/'
traindir
=
datadir
+
'Train/'
testdir
=
datadir
+
'Test/'
#%%
def
get_lr
(
optimizer
):
for
param_group
in
optimizer
.
param_groups
:
print
(
'Learning rate ='
)
print
(
param_group
[
'lr'
])
return
param_group
[
'lr'
]
#To convert data from PIL to tensor
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
=
40
,
shuffle
=
True
,
num_workers
=
0
)
testload
=
datasets
.
ImageFolder
(
root
=
testdir
,
transform
=
transform
)
testset
=
torch
.
utils
.
data
.
DataLoader
(
testload
,
batch_size
=
40
,
shuffle
=
True
,
num_workers
=
0
)
device
=
torch
.
device
(
"cuda:0"
if
torch
.
cuda
.
is_available
()
else
"cpu"
)
#device = torch.device("cpu")
print
(
device
)
net
=
torchvision
.
models
.
resnet18
(
pretrained
=
False
)
#print(net)
#net.to(device)
#summary(net, (3 ,224 ,224))
fc
=
nn
.
Sequential
(
OrderedDict
([
(
'fc1'
,
nn
.
Linear
(
512
,
100
)),
(
'relu'
,
nn
.
ReLU
()),
(
'fc2'
,
nn
.
Linear
(
100
,
4
)),
(
'output'
,
nn
.
LogSoftmax
(
dim
=
1
))
]))
net
.
fc
=
fc
net
.
to
(
device
)
summary
(
net
,
(
3
,
512
,
512
))
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
=
'./ResNet18-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
,
'ResNet18'
)
count_parameters
(
net
)
plotname
=
'ResNet18'
+
'_Stainless_steel'
+
'_confusion_matrix'
+
'.png'
plot_confusion_matrix
(
y_true
,
y_pred
,
classes
,
plotname
)
Event Timeline
Log In to Comment