import torch
from torchvision import d
from torch import nn
from torch.utils.data import dataset, dataloader
from torch import optim
train_loader = torch.utils.data.DataLoader(
torchvision.datasets.MNIST(’/files/’, train=True, download=True,
transform=torchvision.transforms.Compose([
torchvision.transforms.ToTensor(),
torchvision.transforms.Normalize(
(0.1307,), (0.3081,))
])),
batch_size=64, shuffle=True)
val_loader = torch.utils.data.DataLoader(
torchvision.datasets.MNIST(’/files/’, train=False, download=True,
transform=torchvision.transforms.Compose([
torchvision.transforms.ToTensor(),
torchvision.transforms.Normalize(
(0.1307,), (0.3081,))
])),
batch_size=1000, shuffle=True)
model
class CNN(nn.Module):
def init(self):
super(CNN, self).__init__()
self.conv = nn.Sequential(
nn.Conv2d(1, 10, kernel_size=5), # non-RGB image hence input_channel = 1, out_channels is the number of feature maps,
nn.ReLU(),
nn.MaxPool2d(2),
nn.Conv2d(10, 20, kernel_size=5),
nn.ReLU(),
nn.MaxPool2d(2)
)
self.flatten = nn.Flatten()
self.dense = nn.Sequential(
nn.Linear(320, 50),
nn.BatchNorm2d(50),
nn.Linear(50, 10)
)
def forward(self, x):
x = self.conv(x)
x = self.flatten(x)
x = self.dense(x)
return x
model = CNN()
loss_fn = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters())
epochs = 20
train_loss = []
val_loss = []
for t in range(epochs):
for feature, label in train_loader:
# forward
pred = model(feature)
# compute the objective function
output = loss_fn(pred, label)
# clean the gradients
model.zero_grad()
# backward propogation
output.backward()
# update
optimizer.step()
train_loss.append(output.item())
print(f’Epoch {t + 1}, train loss: {output}’)
for feature2, label2 in val_loader:
with torch.no_grad():
# forward
pred2 = model(feature2)
# compute the objective function
output2 = loss_fn(pred2, label2)
val_loss.append(output2.item())
print(f'Epoch {t + 1}, train loss: {output}')
ValueError: expected 4D input (got 2D input)