dcarpintero/mlp-digit-classifier

1import torch2import torch.nn as nn3 4device = "cuda:0" if torch.cuda.is_available() else "cpu"5 6 7class Linear(nn.Module):8    def __init__(self, in_features: int, out_features: int):9        super(Linear, self).__init__()10        self.in_features = in_features11        self.out_features = out_features12 13        self.weight = nn.Parameter(14            (15                torch.randn((self.in_features, self.out_features), device=device) * 0.116            ).requires_grad_()17        )18        self.bias = nn.Parameter(19            (torch.randn(self.out_features, device=device) * 0.1).requires_grad_()20        )21 22    def forward(self, x: torch.Tensor) -> torch.Tensor:23        return x @ self.weight + self.bias24 25 26class ReLU(nn.Module):27    @staticmethod28    def forward(x: torch.Tensor) -> torch.Tensor:29        return torch.max(x, torch.tensor(0))30 31 32class Sequential(nn.Module):33    def __init__(self, *layers):34        super(Sequential, self).__init__()35        self.layers = nn.ModuleList(layers)36 37    def forward(self, x: torch.Tensor) -> torch.Tensor:38        for layer in self.layers:39            x = layer(x)40        return x41 42 43class Flatten(nn.Module):44    @staticmethod45    def forward(x: torch.Tensor) -> torch.Tensor:46        return x.view(x.size(0), -1)47 48 49class DigitClassifier(nn.Module):50    def __init__(self):51        super(DigitClassifier, self).__init__()52        self.main = Sequential(53            Flatten(),54            Linear(in_features=784, out_features=256),55            ReLU(),56            Linear(in_features=256, out_features=64),57            ReLU(),58            Linear(in_features=64, out_features=10),59        )60 61    def forward(self, x: torch.Tensor) -> torch.Tensor:62        return self.main(x)63