forked from pinae/PyTorch-MNIST-Example
-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathtrain.py
114 lines (96 loc) · 3.72 KB
/
train.py
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
import numpy as np
import wandb
import torch
from torch.utils.data import DataLoader
from load_data import *
from MultiLayerPerceptron import FeedForwardNetwork, DeepFeedForwardNetwork
from ConvolutionalNetwork import ConvolutionalNetwork, DeepConvolutionalNetwork
from torch.nn import CrossEntropyLoss
from torch.optim.lr_scheduler import StepLR
def main():
wandb.init(
project="MNIST-Examples",
config={
"dataset": "MNIST",
"learning_rate": 1e-0,
"epochs": 40,
"batch_size": 6000,
"shuffle_data": True
}
)
device = (
"cuda"
if torch.cuda.is_available()
else "mps"
if torch.backends.mps.is_available()
else "cpu"
)
wandb.config["device"] = device
train_dataloader = DataLoader(training_data(),
batch_size=wandb.config["batch_size"],
shuffle=wandb.config["shuffle_data"])
test_dataloader = DataLoader(test_data(),
batch_size=wandb.config["batch_size"],
shuffle=wandb.config["shuffle_data"])
for X, y in test_dataloader:
input_size = X.shape[-1] * X.shape[-2]
print(input_size)
print(f"Shape of X [N, C, H, W]: {X.shape}")
print(f"Shape of y: {y.shape} {y.dtype}")
break
#model = FeedForwardNetwork().to(device)
model = DeepFeedForwardNetwork(input_size=input_size).to(device)
#model = ConvolutionalNetwork().to(device)
#model = DeepConvolutionalNetwork().to(device)
print(model)
loss_fn = CrossEntropyLoss()
wandb.config["loss_fn"] = str(loss_fn)[:-2]
optimizer = torch.optim.SGD(model.parameters(), lr=wandb.config["learning_rate"])
wandb.config["optimizer"] = str(optimizer).split('(')[0].strip()
def train(dataloader, model, loss_fn, optimizer):
size = len(dataloader.dataset)
model.train()
for batch, (X, y) in enumerate(dataloader):
X, y = X.to(device), y.to(device)
# Compute prediction error
pred = model(X)
loss = loss_fn(pred, y)
wandb.log({
"lr": optimizer.state_dict()['param_groups'][0]['lr'],
"train_loss": loss.item(),
})
# Backpropagation
loss.backward()
optimizer.step()
optimizer.zero_grad()
if True or batch % 100 == 0:
loss, current = loss.item(), (batch + 1) * len(X)
print(f"loss: {loss:>7f} [{current:>5d}/{size:>5d}]")
def test(dataloader, model, loss_fn):
size = len(dataloader.dataset)
num_batches = len(dataloader)
model.eval()
test_loss, correct = 0, 0
with torch.no_grad():
for X, y in dataloader:
X, y = X.to(device), y.to(device)
pred = model(X)
test_loss += loss_fn(pred, y).item()
correct += (pred.argmax(1) == y).type(torch.float).sum().item()
test_loss /= num_batches
correct /= size
wandb.log({
"avg_test_loss": test_loss,
"accuracy": 100 * correct
})
print(f"Test Error: \n Accuracy: {(100*correct):>0.1f}%, Avg loss: {test_loss:>8f} \n")
scheduler = StepLR(optimizer, step_size=2, gamma=0.75)
for t in range(wandb.config["epochs"]):
print(f"Epoch {t+1} (lr: {optimizer.state_dict()['param_groups'][0]['lr']})\n-------------------------------")
train(train_dataloader, model, loss_fn, optimizer)
test(test_dataloader, model, loss_fn)
scheduler.step()
wandb.finish()
if __name__ == "__main__":
for i in range(10):
main()