import torch
import torch.nn as nn

class MiniVisionV2(nn.Module):
    def __init__(self):
        super().__init__()

        self.model = nn.Sequential(
            nn.Conv2d(1, 32, 3, padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.MaxPool2d(2, 2),

            nn.Conv2d(32, 64, 3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.MaxPool2d(2, 2),

            nn.Flatten(),
            nn.Linear(3136, 256),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(256, 10)
        )

    def forward(self, x):
        x = self.model(x)
        return x

if __name__ == '__main__':
    minivisionv2 = MiniVisionV2()
    params = sum(p.numel() for p in minivisionv2.parameters())
    print(f"Total params: {params / 1000000:,}M")

    input = torch.randn(64, 1, 28, 28)
    with torch.no_grad():
        output = minivisionv2(input)
    print(output)