Learn practical skills, build real-world projects, and advance your career
import torch
import torchvision
from torchvision import transforms, datasets
train = datasets.MNIST('', train=True, download=True, transform=transforms.Compose([transforms.ToTensor()]))
test = datasets.MNIST('', train=False, download=True, transform=transforms.Compose([transforms.ToTensor()]))
trainset = torch.utils.data.DataLoader(train, batch_size = 10, shuffle=True)
testset = torch.utils.data.DataLoader(test, batch_size = 10, shuffle=False)
import torch.nn as nn
import torch.nn.functional as F
class Net(nn.Module): #inherit from nn.module
    def __init__(self): 
        #initialize nn.module
        super().__init__()
        #define fully connected layers to nn
        self.fc1 = nn.Linear(28*28, 64)
        self.fc2 = nn.Linear(64, 64)
        self.fc3 = nn.Linear(64, 64)
        self.fc4 = nn.Linear(64, 10) #we have 10 classes
        
    def forward(self,x):
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = F.relu(self.fc3(x))
        x = self.fc4(x) #we dont want to pass relu here. because this is o/p layer.
        return F.log_softmax(x, dim=1) #dim=1 is similar to axes, which this is the probability distribution we want to sum to 1.
        return x
    
net = Net()
net
Net(
  (fc1): Linear(in_features=784, out_features=64, bias=True)
  (fc2): Linear(in_features=64, out_features=64, bias=True)
  (fc3): Linear(in_features=64, out_features=64, bias=True)
  (fc4): Linear(in_features=64, out_features=10, bias=True)
)