File size: 2,900 Bytes
9d8daba |
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 |
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
import scipy
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")# assign the device to the model
#hyper-parameters
learning_rate = 0.005
batch_size = 128
hidden_size = 300
num_classes = 10
num_epochs = 550
#load data
transform = transforms.Compose(
[transforms.Resize(64), transforms.CenterCrop(64), transforms.ToTensor()]
)
train_dataset = torchvision.datasets.STL10(root='./dataSTL10', split="train", transform=transform, download=True)
test_dataset = torchvision.datasets.STL10(root='./dataSTL10', split="test", transform=transform, download=True)
train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True)
test_loader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=batch_size, shuffle=False)
# CNN
class CNN(nn.Module):
def __init__(self):
super(CNN, self).__init__()
self.conv1 = nn.Conv2d(3, 32, 5)
self.conv2 = nn.Conv2d(32, 64, 5)
#full layer
self.fc1 = nn.Linear(64 * 13 * 13, 128)
self.fc2 = nn.Linear(128, 64)
self.fc3 = nn.Linear(64, num_classes)
def forward(self, x):
x = F.max_pool2d(F.relu(self.conv1(x)), (2,2))
x = F.max_pool2d(F.relu(self.conv2(x)), 2)
x = x.view(-1, self.num_flat_features(x))
x = F.relu(self.fc1(x))
x = F.relu(self.fc2(x))
x = self.fc3(x)
return x
def num_flat_features(self, x):
size = x.size()[1:] # all dimensions except the batch dimension
num_features = 1
for s in size:
num_features *= s
return num_features
cnn = CNN().to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(cnn.parameters(), lr=learning_rate)
# training loop
for epoch in range(num_epochs):
for i, (images, labels) in enumerate(train_loader):
images = images.to(device)
labels = labels.to(device)
out = cnn(images)
loss = criterion(out, labels)
optimizer.zero_grad()
loss.backward()
optimizer.step()
if(i+1) % 1 == 0:
print(f'epoch: {epoch+1}/{num_epochs} step: {i+1}, loss: loss: {loss.item():.4f}')
with torch.no_grad():
n_correct = 0
n_samples = 0
for images, labels in test_loader:
images = images.to(device)
labels = labels.to(device)
outputs = cnn(images)
# max returns (value ,index)
_, predicted = torch.max(outputs.data, 1)
n_samples += labels.size(0)
n_correct += (predicted == labels).sum().item()
acc = 100.0 * n_correct / n_samples
print(f'Accuracy of the network on the {n_samples} test images: {acc} %')
# Save the model
torch.save(cnn.state_dict(), "cnn_model.pth")
|