81 lines
2.8 KiB
Python
81 lines
2.8 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
import torch.optim as optim
|
|
import torch.nn.functional as F
|
|
from torchvision import transforms, datasets
|
|
from tqdm import tqdm
|
|
import wandb
|
|
|
|
wandb.init(project='pytorch_ai')
|
|
train = datasets.MNIST('./datasets', train=True, download=True,
|
|
transform=transforms.Compose([
|
|
transforms.ToTensor()
|
|
]))
|
|
|
|
test = datasets.MNIST('./datasets', train=False, download=True,
|
|
transform=transforms.Compose([
|
|
transforms.ToTensor()
|
|
]))
|
|
|
|
trainset = torch.utils.data.DataLoader(train, batch_size=200, shuffle=True)
|
|
testset = torch.utils.data.DataLoader(test, batch_size=10, shuffle=False)
|
|
|
|
|
|
class Net(nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.fc1 = nn.Linear(28 * 28, 64)
|
|
self.fc2 = nn.Linear(64, 120)
|
|
self.fc3 = nn.Linear(120, 120)
|
|
self.fc4 = nn.Linear(120, 64)
|
|
self.fc5 = nn.Linear(64, 10)
|
|
|
|
def forward(self, x):
|
|
x = F.relu(self.fc1(x))
|
|
x = F.relu(self.fc2(x))
|
|
x = F.relu(self.fc3(x))
|
|
x = F.relu(self.fc4(x))
|
|
x = self.fc5(x)
|
|
return F.log_softmax(x, dim=1)
|
|
|
|
|
|
net = Net()
|
|
wandb.watch(net)
|
|
|
|
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
|
|
print('runnning on %s' % device)
|
|
|
|
net = net.to(device)
|
|
|
|
loss_function = nn.CrossEntropyLoss()
|
|
optimizer = optim.Adam(net.parameters(), lr=0.001)
|
|
|
|
for epoch in range(200): # 10 full passes over the data
|
|
for data in tqdm(trainset): # `data` is a batch of data
|
|
X, y = data # X is the batch of features, y is the batch of targets.
|
|
net.zero_grad() # sets gradients to 0 before loss calc. You will do this likely every step.
|
|
X = X.to(device)
|
|
output = net(X.view(-1, 784)) # pass in the reshaped batch (recall they are 28x28 atm)
|
|
output = output.cpu()
|
|
loss = loss_function(output, y) # calc and grab the loss value
|
|
loss.backward() # apply this loss backwards thru the network's parameters
|
|
optimizer.step() # attempt to optimize weights to account for loss/gradients
|
|
wandb.log({'loss': loss})
|
|
net = net.cpu()
|
|
torch.save(net, './nets/net_gpu_large_batch_' + str(epoch) + ".pt")
|
|
net = net.to(device)
|
|
correct = 0
|
|
total = 0
|
|
with torch.no_grad():
|
|
for data in testset:
|
|
X, y = data
|
|
X = X.to(device)
|
|
output = net(X.view(-1, 784))
|
|
output = output.cpu()
|
|
for idx, i in enumerate(output):
|
|
if torch.argmax(i) == y[idx]:
|
|
correct += 1
|
|
total += 1
|
|
wandb.log({'test_accuracy': correct / total})
|
|
print("Accuracy: ", round(correct / total, 3))
|
|
wandb.log({'epoch': epoch})
|