import torch
import torchvision
import torchvision.transforms as transforms
transform = transforms.Compose([transforms.ToTensor()])
trainset = torchvision.datasets.MNIST(root='.data/', train=True, download=True, transform=transform)
train_data = torch.utils.data.DataLoader(trainset, shuffle=True, batch_size=4,num_workers=2)
testset = torchvision.datasets.MNIST(root='.data/', train=False,download=True, transform=transform )
test_data = torch.utils.data.DataLoader(testset, shuffle=True, batch_size=4)
class SimpleConv(torch.nn.Module):
def __init__(self):
super(SimpleConv, self).__init__()
self.conv1 = torch.nn.Conv2d(in_channels=1, out_channels=6, kernel_size=5)
self.pool = torch.nn.MaxPool2d(kernel_size=2, stride=2)
#return 28/2 = 14*14
self.layer1 = torch.nn.Sequential(
# 1 - count of input map
# 32 - count of output map
torch.nn.Conv2d(1,32, kernel_size=5, stride=1),
torch.nn.ReLU(),
torch.nn.MaxPool2d(kernel_size=2, stride=2)
)
#return 14/2 = 7*7
self.layer2 = torch.nn.Sequential(
torch.nn.Conv2d(32,64,kernel_size=5,stride=1,padding=2),
torch.nn.ReLU(),
torch.nn.MaxPool2d(kernel_size=2,stride=2)
)
self.drop_out = torch.nn.Dropout()
self.fc1 = torch.nn.Linear(7 * 7 * 64, 1000)
self.fc2 = torch.nn.Linear(1000, 10)
def forward(self, x):
out = self.layer1(x)
out = self.layer2(out)
print(out.shape)
out = x.view(-1, 7 * 7 * 64)
print(out.shape)
out = self.drop_out(out)
out = self.fc1(out)
out = self.fc2(out)
return out
model = SimpleConv()
loss_fun = torch.nn.CrossEntropyLoss()
optim = torch.optim.Adam(model.parameters(),lr= 0.001)
num_epochs = 5
num_classes = 10
for epoch in range(num_epochs):
for i, batch in enumerate(train_data):
X_batch, y_batch = batch
print(X_batch.shape)
print(y_batch.shape)
optim.zero_grad()
output = model(X_batch)
loss = loss_fun(output, y_batch)
Your out layer is of shape [4,64,6,6] so you should reshape first out not x and secondly, you should hence reshape by [-1,64 x 6 x 6]. see below
...
self.fc1 = torch.nn.Linear(64*6*6, 1000)
self.fc2 = torch.nn.Linear(1000, 10)
def forward(self, x):
out = self.layer1(x)
out = self.layer2(out)
print(out.shape)
out = out.view(-1, 64*6*6)
print(out.shape)
out = self.drop_out(out)
out = self.fc1(out)
out = self.fc2(out)
return out