Update network model structure:

Changed model from FCNN to CNN
This commit is contained in:
Cezary Adamczak 2021-06-20 15:00:34 +02:00
parent 0b2791f961
commit 3898a3bcab

View File

@ -1,69 +1,76 @@
import PIL
import numpy as np
import torch import torch
import torch.nn as nn
import torch.optim as optim
import torchvision import torchvision
import torchvision.transforms as transforms import torchvision.transforms as transforms
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import numpy as np
from matplotlib.pyplot import imshow
import os
import PIL
import numpy as np
from matplotlib.pyplot import imshow from matplotlib.pyplot import imshow
def to_negative(img): def to_negative(img):
img = PIL.ImageOps.invert(img) img = PIL.ImageOps.invert(img)
return img return img
class Negative(object): class Negative(object):
def __init__(self): def __init__(self):
pass pass
def __call__(self, img): def __call__(self, img):
return to_negative(img) return to_negative(img)
def plotdigit(image):
img = np.reshape(image, (-1, 100))
imshow(img, cmap='Greys')
transform = transforms.Compose([Negative(), transforms.ToTensor()]) transform = transforms.Compose([Negative(), transforms.ToTensor()])
train_set = torchvision.datasets.ImageFolder(root='../src/train', transform=transform) train_set = torchvision.datasets.ImageFolder(root='train', transform=transform)
classes = ("apple", "potato") classes = ("pepper", "potato", "strawberry", "tomato")
BATCH_SIZE = 2 BATCH_SIZE = 4
train_loader = torch.utils.data.DataLoader(train_set, batch_size=BATCH_SIZE, shuffle=True, num_workers=0) train_loader = torch.utils.data.DataLoader(train_set, batch_size=BATCH_SIZE, shuffle=True, num_workers=0)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
class Net(nn.Module): class Net(nn.Module):
def __init__(self): def __init__(self):
super(Net, self).__init__() super().__init__()
self.flatten = nn.Flatten() self.network = nn.Sequential(
self.linear_relu_stack = nn.Sequential( nn.Conv2d(3, 32, kernel_size=3, padding=1), #3 channels to 32 channels
nn.Linear(3 * 100 * 100, 512),
nn.ReLU(), nn.ReLU(),
nn.Linear(512, 512), nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1),
nn.ReLU(), nn.ReLU(),
nn.Linear(512, 2), nn.MaxPool2d(2, 2), # output: 64 channels x 50 x 50 image size - decrease
nn.ReLU()
)
self.linear_relu_stack = self.linear_relu_stack.to(device)
def forward(self, x): nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1),
x = self.flatten(x).to(device) nn.ReLU(),
logits = self.linear_relu_stack(x).to(device) nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1), #increase power of model
return logits nn.ReLU(),
nn.MaxPool2d(2, 2), # output: 128 x 25 x 25
nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1),
nn.ReLU(),
nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1),
nn.ReLU(),
nn.MaxPool2d(5, 5), # output: 256 x 5 x 5
nn.Flatten(), #a single vector 256*5*5,
nn.Linear(256*5*5, 1024),
nn.ReLU(),
nn.Linear(1024, 512),
nn.ReLU(),
nn.Linear(512, 4))
def forward(self, xb):
return self.network(xb)
def training_network(): def training_network():
net = Net() net = Net()
net = net.to(device) net = net.to(device)
criterion = nn.CrossEntropyLoss() criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9) optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)
for epoch in range(4): for epoch in range(10):
running_loss = 0.0 running_loss = 0.0
for i, data in enumerate(train_loader, 0): for i, data in enumerate(train_loader, 0):
inputs, labels = data[0].to(device), data[1].to(device) inputs, labels = data[0].to(device), data[1].to(device)
@ -72,34 +79,33 @@ def training_network():
loss = criterion(outputs, labels) loss = criterion(outputs, labels)
loss.backward() loss.backward()
optimizer.step() optimizer.step()
running_loss += loss.item() running_loss += loss.item()
if i % 2000 == 1999: if i % 200 == 199:
print('[%d, %5d] loss: %.3f' % (epoch + 1, i + 1, running_loss)) print('[%d, %5d] loss: %.3f' % (epoch + 1, i + 1, running_loss))
running_loss = 0.0 running_loss = 0.0
print("Finished training") print("Finished training")
save_network_to_file(net) save_network_to_file(net)
def result_from_network(net, loaded_image): def result_from_network(net, loaded_image):
image = PIL.Image.open(loaded_image) image = PIL.Image.open(loaded_image)
pil_to_tensor = transforms.ToTensor()(image.convert("RGB")).unsqueeze_(0) pil_to_tensor = transforms.Compose([Negative(), transforms.ToTensor()])(image.convert("RGB")).unsqueeze_(0)
outputs = net(pil_to_tensor.to(device)) outputs = net(pil_to_tensor)
return classes[torch.max(outputs, 1)[1]] return classes[torch.max(outputs, 1)[1]]
def save_network_to_file(network): def save_network_to_file(network):
torch.save(network.state_dict(), 'network_model.pth') torch.save(network.state_dict(), 'network_model.pth')
print("Network saved to file") print("Network saved to file")
def load_network_from_structure(network): def load_network_from_structure(network):
network.load_state_dict(torch.load('network_model.pth')) network.load_state_dict(torch.load('network_model.pth'))
# Create network_model.pth
if __name__ == "__main__": if __name__ == "__main__":
print(torch.cuda.is_available()) print(torch.cuda.is_available())
training_network() training_network()