Spaces:
Build error
Build error
File size: 3,784 Bytes
583e46a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 | """
Custom CNN model built from scratch.
5 convolutional blocks with batch norm + global average pooling.
"""
import torch
import torch.nn as nn
class ConvBlock(nn.Module):
"""Conv2d -> BatchNorm -> ReLU -> MaxPool"""
def __init__(self, in_channels: int, out_channels: int, kernel_size: int = 3,
padding: int = 1, pool_size: int = 2):
super().__init__()
self.block = nn.Sequential(
nn.Conv2d(in_channels, out_channels, kernel_size=kernel_size, padding=padding),
nn.BatchNorm2d(out_channels),
nn.ReLU(inplace=True),
nn.Conv2d(out_channels, out_channels, kernel_size=kernel_size, padding=padding),
nn.BatchNorm2d(out_channels),
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=pool_size, stride=pool_size),
)
def forward(self, x: torch.Tensor) -> torch.Tensor:
return self.block(x)
class CattleCNN(nn.Module):
"""
Custom CNN with 5 convolutional blocks.
Each block: 2x(Conv2d + BN + ReLU) + MaxPool
Ends with Global Average Pooling + FC classifier.
"""
def __init__(
self,
num_classes: int = 26,
in_channels: int = 3,
conv_channels: list[int] = None,
kernel_size: int = 3,
padding: int = 1,
pool_size: int = 2,
dropout: float = 0.4,
use_global_avg_pool: bool = True,
):
super().__init__()
if conv_channels is None:
conv_channels = [32, 64, 128, 256, 512]
self.use_global_avg_pool = use_global_avg_pool
# Build conv blocks
blocks = []
prev_channels = in_channels
for out_ch in conv_channels:
blocks.append(ConvBlock(prev_channels, out_ch, kernel_size, padding, pool_size))
prev_channels = out_ch
self.features = nn.Sequential(*blocks)
# Global average pooling
self.gap = nn.AdaptiveAvgPool2d(1)
# Classifier head
self.classifier = nn.Sequential(
nn.Dropout(p=dropout),
nn.Linear(conv_channels[-1], 256),
nn.BatchNorm1d(256),
nn.ReLU(inplace=True),
nn.Dropout(p=dropout * 0.5),
nn.Linear(256, num_classes),
)
self._init_weights()
def _init_weights(self):
for m in self.modules():
if isinstance(m, nn.Conv2d):
nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
if m.bias is not None:
nn.init.constant_(m.bias, 0)
elif isinstance(m, (nn.BatchNorm2d, nn.BatchNorm1d)):
nn.init.constant_(m.weight, 1)
nn.init.constant_(m.bias, 0)
elif isinstance(m, nn.Linear):
nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
nn.init.constant_(m.bias, 0)
def forward(self, x: torch.Tensor) -> torch.Tensor:
x = self.features(x)
x = self.gap(x)
x = x.view(x.size(0), -1)
x = self.classifier(x)
return x
@staticmethod
def from_config(config: dict) -> 'CattleCNN':
"""Create model from config dict."""
arch = config.get('model', {}).get('architecture', {})
return CattleCNN(
num_classes=config.get('num_classes', 26),
in_channels=config.get('image', {}).get('channels', 3),
conv_channels=arch.get('conv_channels', [32, 64, 128, 256, 512]),
kernel_size=arch.get('kernel_size', 3),
padding=arch.get('padding', 1),
pool_size=arch.get('pool_size', 2),
dropout=arch.get('dropout', 0.4),
use_global_avg_pool=arch.get('use_global_avg_pool', True),
)
|