Architectures / Computer Vision

🖼️ Simple CNN

Simple Convolutional Neural Network for image classification

Every number on this page is computed from the graph by the same functions the app runs, not written by hand.

When to pick it

Structure

9 layers. Output shapes are propagated from the input shape, batch dimension excluded.

What the verifier says

The same 41 structural checks that run on every edit in the app, on this graph.

The PyTorch it exports

Generated from the graph above.

# Architecture designed with Neurarch: https://neurarch.com
# PyTorch: compatible with Python 3.8+ and torch>=1.12
# Colab: pip install torch torchvision  (usually pre-installed)

import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Tuple

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

        self.conv2d_1 = nn.Conv2d(1, 32, kernel_size=3, stride=1, padding=1)
        self.maxpool2d_1 = nn.MaxPool2d(kernel_size=2, stride=2, padding=0)
        self.linear_1 = nn.Linear(6272, 128)
        self.linear_2 = nn.Linear(128, 10)

    def forward(self, x):
        # Input shape: [1,28,28]
        conv2d_nv2d_1 = self.conv2d_1(x)
        relu_relu_1 = F.relu(conv2d_nv2d_1)
        maxpool2d_ol2d_1 = self.maxpool2d_1(relu_relu_1)
        flatten_tten_1 = torch.flatten(maxpool2d_ol2d_1, 1)
        linear_near_1 = self.linear_1(flatten_tten_1)
        relu_relu_2 = F.relu(linear_near_1)
        linear_near_2 = self.linear_2(relu_relu_2)
        # Output
        return linear_near_2


if __name__ == '__main__':
    model = SimpleCNN()
    model.eval()

    x = torch.randn(1, 28, 28)  # (batch, seq_len, embed_dim)
    with torch.no_grad():
        output = model(x)

    print(f'Input  shape : {tuple(x.shape)}')
    print(f'Output shape : {tuple(output.shape)}')
    total = sum(p.numel() for p in model.parameters())
    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    print(f'Parameters   : {total:,} total, {trainable:,} trainable')

For agents

Also in Computer Vision