import torch.nn as nn

# 2D convolution (for images)
conv = nn.Conv2d(
    in_channels=3,    # Input channels (RGB=3, grayscale=1)
    out_channels=16,  # Number of filters
    kernel_size=3,    # 3x3 filter
    stride=1,         # Filter movement step
    padding=1         # Preserves spatial dimensions (with stride=1)
)

# Shape transformation: (batch, channels, height, width)
x = torch.randn(32, 3, 64, 64)  # 32 RGB images of 64x64
print(conv(x).shape)  # → torch.Size([32, 16, 64, 64])

# Max pooling (common for downsampling)
pool = nn.MaxPool2d(kernel_size=2, stride=2)
print(pool(conv(x)).shape)  # → torch.Size([32, 16, 32, 32])

# Adaptive pooling (useful for varying input sizes)
adaptive_pool = nn.AdaptiveAvgPool2d((7, 7))
print(adaptive_pool(x).shape)  # → torch.Size([32, 3, 7, 7])

# Batch Normalization
bn = nn.BatchNorm2d(16)  # num_features = out_channels
x = conv(x)
x = bn(x)

# Layer Normalization (for NLP/sequences)
ln = nn.LayerNorm([16, 64, 64])

# Spatial dropout (drops entire channels)
dropout = nn.Dropout2d(p=0.25)

class CNN(nn.Module):
    def __init__(self, num_classes=10):
        super().__init__()
        self.features = nn.Sequential(
            # Block 1
            nn.Conv2d(3, 32, kernel_size=3, padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.MaxPool2d(2),
            
            # Block 2
            nn.Conv2d(32, 64, kernel_size=3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.MaxPool2d(2),
            
            # Block 3
            nn.Conv2d(64, 128, kernel_size=3, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(),
            nn.MaxPool2d(2),
        )
        
        self.classifier = nn.Sequential(
            nn.Linear(128 * 4 * 4, 512),  # Adjusted based on input size
            nn.ReLU(),
            nn.Dropout(0.5),
            nn.Linear(512, num_classes)
        )
        
    def forward(self, x):
        x = self.features(x)
        x = torch.flatten(x, 1)  # Flatten all dimensions except batch
        x = self.classifier(x)
        return x

# Usage
model = CNN().to(device)
print(model)

Hₒᵤₜ = ⌊(Hᵢₙ + 2P - K)/S⌋ + 1
Wₒᵤₜ = ⌊(Wᵢₙ + 2P - K)/S⌋ + 1

#VisionTransformer #ViT #DeepLearning #ComputerVision #Transformers #AI #MachineLearning #NeuralNetworks #ImageClassification #AttentionIsAllYouNeed

✉️ Our Telegram channels: https://yangx.top/addlist/0f6vfFbEMdAwODBk

📱 Our WhatsApp channel: https://whatsapp.com/channel/0029VaC7Weq29753hpcggW2A

#VisionTransformer #ViTFromScratch #PyTorch #DeepLearning #ComputerVision #Transformers #AI #MachineLearning #CodingTutorial #AttentionIsAllYouNeed

✉️ Our Telegram channels: https://yangx.top/addlist/0f6vfFbEMdAwODBk

📱 Our WhatsApp channel: https://whatsapp.com/channel/0029VaC7Weq29753hpcggW2A

#VisionTransformer #TransferLearning #HuggingFace #ImageNet #FineTuning #AI #DeepLearning #ComputerVision #Transformers #ModelZoo