File size: 2,861 Bytes
a74b1a9 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 |
#!/usr/bin/env python
# coding: utf-8
# In[2]:
# import pytorch and machine learning stuff
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import DataLoader
# import other stuff
import numpy as np
# In[3]:
# import resnet from torch
import torch.library
from torchvision.models import squeezenet1_1
from torchvision.models import resnet50
from torchvision.models import resnet18
from torchvision.models import mobilenet_v2
from torchvision import transforms
from torchvision.datasets import ImageFolder
from torch.utils.data import DataLoader
# In[4]:
class_num = 5
classes = ['Ak', 'Ala_Idris', 'Buzgulu', 'Dimnit', 'Nazli']
# In[5]:
model = mobilenet_v2(pretrained=True)
# In[9]:
print(model)
# In[7]:
transform = transforms.Compose([
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize([0.485,0.456,0.406],[0.229,0.224,0.225])
])
# In[8]:
training_set = ImageFolder('../data/train', transform=transform)
test_set = ImageFolder('../data/test', transform=transform)
val_set = ImageFolder('../data/val', transform=transform)
# In[47]:
batch_size = 8
epochs = 5
lr = 1e-5
loss_fn = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=lr)
# In[48]:
train_loader = DataLoader(training_set, batch_size=batch_size, shuffle=True)
test_loader = DataLoader(test_set, batch_size=batch_size)
val_loader = DataLoader(val_set, batch_size=batch_size)
# In[49]:
model.classifier[1] = nn.Linear(in_features=1280, out_features=class_num)
# In[52]:
epochs = 1
# In[55]:
# train the model
for epoch in range(epochs):
model.train()
for batch_idx, (data, target) in enumerate(train_loader):
optimizer.zero_grad()
output = model(data)
print("Out: ", [a.argmax().item() for a in output])
print("Target: ", target)
loss = loss_fn(output, target)
loss.backward()
optimizer.step()
if batch_idx % 10 == 0:
print('Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
epoch, batch_idx * len(data), len(train_loader.dataset),
100. * batch_idx / len(train_loader), loss.item()
))
# test the model
model.eval()
test_loss = 0
correct = 0
with torch.no_grad():
for data, target in test_loader:
output = model(data)
test_loss += loss_fn(output, target).item()
pred = output.argmax(dim=1, keepdim=True)
correct += pred.eq(target.view_as(pred)).sum().item()
test_loss /= len(test_loader.dataset)
print('\nTest set: Avg. loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
test_loss, correct, len(test_loader.dataset),
100. * correct / len(test_loader.dataset)
))
# In[51]:
model_scripted = torch.jit.script(model)
model_scripted.save('../models/mobilenet.pt')
# In[ ]:
|