Update network model structure:
Changed model from FCNN to CNN
This commit is contained in:
parent
0b2791f961
commit
3898a3bcab
@ -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()
|
Loading…
Reference in New Issue
Block a user