Modern CNN Architectures

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

# Compound Scaling Coefficients for EfficientNet
# Determined via grid search on EfficientNet-B0
 
# Base scaling coefficients
ALPHA = 1.2   # Depth scaling base
BETA = 1.1    # Width scaling base  
GAMMA = 1.15  # Resolution scaling base
 
# Verify constraint: α · β² · γ² ≈ 2
constraint = ALPHA * (BETA ** 2) * (GAMMA ** 2)
print(f"Constraint value: {constraint:.3f}")  # Should be ≈ 2.0
 
def compute_scaling(phi: float) -> dict:
    """
    Compute scaling factors for a given compound coefficient phi.
    
    Args:
        phi: Compound scaling coefficient (e.g., 0, 1, 2, ...)
        
    Returns:
        Dictionary with depth, width, and resolution multipliers
    """
    depth_multiplier = ALPHA ** phi
    width_multiplier = BETA ** phi
    resolution_multiplier = GAMMA ** phi
    
    # Approximate FLOPs scaling
    flops_scaling = (ALPHA * (BETA ** 2) * (GAMMA ** 2)) ** phi
    
    return {
        "depth": depth_multiplier,
        "width": width_multiplier, 
        "resolution": resolution_multiplier,
        "approx_flops": flops_scaling
    }
 
# EfficientNet family scaling
for phi in range(8):  # B0 through B7
    scaling = compute_scaling(phi)
    print(f"EfficientNet-B{phi}: depth={scaling['depth']:.2f}x, "
          f"width={scaling['width']:.2f}x, res={scaling['resolution']:.2f}x, "
          f"~{scaling['approx_flops']:.1f}x FLOPs")

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

import torch
import torch.nn as nn
 
class SqueezeExcitation(nn.Module):
    """
    Squeeze-and-Excitation block for channel attention.
    
    Computes channel-wise importance weights by:
    1. Global average pooling to get channel statistics
    2. Two FC layers to compute attention weights
    3. Sigmoid activation to get scaling factors
    """
    def __init__(self, in_channels: int, reduced_dim: int):
        super().__init__()
        self.se = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),  # Global pooling: (B, C, H, W) -> (B, C, 1, 1)
            nn.Conv2d(in_channels, reduced_dim, 1),  # Reduction
            nn.SiLU(inplace=True),  # Swish activation
            nn.Conv2d(reduced_dim, in_channels, 1),  # Expansion
            nn.Sigmoid()  # Attention weights in [0, 1]
        )
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return x * self.se(x)  # Channel-wise scaling
 
 
class MBConv(nn.Module):
    """
    Mobile Inverted Bottleneck Convolution (MBConv) block.
    
    Structure:
    Input -> Expand (1x1) -> Depthwise (kxk) -> SE -> Project (1x1) -> Output
         \---------------------------------------------------/
                            (residual if in_ch == out_ch)
    
    Args:
        in_channels: Input channel count
        out_channels: Output channel count
        kernel_size: Spatial kernel size for depthwise conv
        stride: Stride for depthwise conv (for downsampling)
        expand_ratio: Channel expansion factor (1 for MBConv1, 6 for MBConv6)
        se_ratio: Squeeze-Excitation reduction ratio (typically 0.25)
    """
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: int = 3,
        stride: int = 1,
        expand_ratio: int = 6,
        se_ratio: float = 0.25
    ):
        super().__init__()
        self.stride = stride
        self.use_residual = (stride == 1) and (in_channels == out_channels)
        
        # Expanded channel dimension
        expanded_channels = in_channels * expand_ratio
        
        layers = []
        
        # 1. Expansion phase (skip if expand_ratio == 1)
        if expand_ratio != 1:
            layers.extend([
                nn.Conv2d(in_channels, expanded_channels, 1, bias=False),
                nn.BatchNorm2d(expanded_channels),
                nn.SiLU(inplace=True)  # Swish activation
            ])
        
        # 2. Depthwise convolution
        layers.extend([
            nn.Conv2d(
                expanded_channels, expanded_channels, kernel_size,
                stride=stride, padding=kernel_size // 2,
                groups=expanded_channels, bias=False  # groups=channels for depthwise
            ),
            nn.BatchNorm2d(expanded_channels),
            nn.SiLU(inplace=True)
        ])
        
        # 3. Squeeze-and-Excitation
        reduced_dim = max(1, int(in_channels * se_ratio))
        layers.append(SqueezeExcitation(expanded_channels, reduced_dim))
        
        # 4. Projection (linear, no activation)
        layers.extend([
            nn.Conv2d(expanded_channels, out_channels, 1, bias=False),
            nn.BatchNorm2d(out_channels)
        ])
        
        self.block = nn.Sequential(*layers)
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        out = self.block(x)
        if self.use_residual:
            out = out + x  # Residual connection
        return out

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
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159

import torch
import torch.nn as nn
from typing import List, Tuple
from dataclasses import dataclass
 
@dataclass
class EfficientNetConfig:
    """Configuration for EfficientNet model family."""
    # (expand_ratio, channels, layers, kernel, stride)
    stage_configs: List[Tuple[int, int, int, int, int]] = None
    
    # Compound scaling parameters
    width_mult: float = 1.0
    depth_mult: float = 1.0
    resolution: int = 224
    dropout_rate: float = 0.2
    
    def __post_init__(self):
        if self.stage_configs is None:
            # EfficientNet-B0 baseline configuration
            # (expand_ratio, out_channels, num_layers, kernel_size, stride)
            self.stage_configs = [
                (1, 16, 1, 3, 1),   # Stage 2: MBConv1
                (6, 24, 2, 3, 2),   # Stage 3: MBConv6
                (6, 40, 2, 5, 2),   # Stage 4: MBConv6
                (6, 80, 3, 3, 2),   # Stage 5: MBConv6
                (6, 112, 3, 5, 1),  # Stage 6: MBConv6
                (6, 192, 4, 5, 2),  # Stage 7: MBConv6
                (6, 320, 1, 3, 1),  # Stage 8: MBConv6
            ]
 
# Compound scaling coefficients
EFFICIENTNET_PARAMS = {
    # (phi, resolution, dropout)
    'b0': (0, 224, 0.2),
    'b1': (0.5, 240, 0.2),
    'b2': (1, 260, 0.3),
    'b3': (2, 300, 0.3),
    'b4': (3, 380, 0.4),
    'b5': (4, 456, 0.4),
    'b6': (5, 528, 0.5),
    'b7': (6, 600, 0.5),
}
 
def get_efficientnet_config(model_name: str) -> EfficientNetConfig:
    """
    Get configuration for a specific EfficientNet variant.
    
    Applies compound scaling coefficients to the B0 baseline.
    """
    if model_name not in EFFICIENTNET_PARAMS:
        raise ValueError(f"Unknown model: {model_name}")
    
    phi, resolution, dropout = EFFICIENTNET_PARAMS[model_name]
    
    # Compute scaling multipliers
    # α = 1.2, β = 1.1, γ = 1.15
    alpha, beta, gamma = 1.2, 1.1, 1.15
    
    depth_mult = alpha ** phi
    width_mult = beta ** phi
    
    return EfficientNetConfig(
        width_mult=width_mult,
        depth_mult=depth_mult,
        resolution=resolution,
        dropout_rate=dropout
    )
 
 
class EfficientNet(nn.Module):
    """
    EfficientNet model implementation.
    
    Builds the network from configuration, applying compound scaling
    to depth (layer counts) and width (channel counts).
    """
    def __init__(self, config: EfficientNetConfig, num_classes: int = 1000):
        super().__init__()
        self.config = config
        
        # Scale helper functions
        def scale_depth(d: int) -> int:
            return max(1, int(d * config.depth_mult + 0.5))
        
        def scale_width(w: int) -> int:
            # Round to nearest multiple of 8 for efficiency
            w = int(w * config.width_mult + 0.5)
            return max(8, (w + 4) // 8 * 8)
        
        # Stem: Initial conv layer
        stem_channels = scale_width(32)
        self.stem = nn.Sequential(
            nn.Conv2d(3, stem_channels, 3, stride=2, padding=1, bias=False),
            nn.BatchNorm2d(stem_channels),
            nn.SiLU(inplace=True)
        )
        
        # Build MBConv stages
        stages = []
        in_channels = stem_channels
        
        for expand, out_ch, layers, kernel, stride in config.stage_configs:
            out_channels = scale_width(out_ch)
            num_layers = scale_depth(layers)
            
            stage = self._build_stage(
                in_channels, out_channels, num_layers,
                expand, kernel, stride
            )
            stages.append(stage)
            in_channels = out_channels
        
        self.stages = nn.Sequential(*stages)
        
        # Head: Final conv, pooling, and classifier
        head_channels = scale_width(1280)
        self.head = nn.Sequential(
            nn.Conv2d(in_channels, head_channels, 1, bias=False),
            nn.BatchNorm2d(head_channels),
            nn.SiLU(inplace=True),
            nn.AdaptiveAvgPool2d(1),
            nn.Flatten(),
            nn.Dropout(config.dropout_rate),
            nn.Linear(head_channels, num_classes)
        )
    
    def _build_stage(
        self, in_ch: int, out_ch: int, num_layers: int,
        expand: int, kernel: int, stride: int
    ) -> nn.Sequential:
        """Build a stage of MBConv blocks."""
        layers = []
        for i in range(num_layers):
            layers.append(MBConv(
                in_ch if i == 0 else out_ch,
                out_ch,
                kernel_size=kernel,
                stride=stride if i == 0 else 1,  # Only first layer can downsample
                expand_ratio=expand
            ))
        return nn.Sequential(*layers)
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.stem(x)
        x = self.stages(x)
        x = self.head(x)
        return x
 
 
# Factory function
def efficientnet_b0(num_classes: int = 1000) -> EfficientNet:
    return EfficientNet(get_efficientnet_config('b0'), num_classes)
 
def efficientnet_b3(num_classes: int = 1000) -> EfficientNet:
    return EfficientNet(get_efficientnet_config('b3'), num_classes)
 
def efficientnet_b7(num_classes: int = 1000) -> EfficientNet:
    return EfficientNet(get_efficientnet_config('b7'), num_classes)

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

import torch
import torch.nn as nn
 
class FusedMBConv(nn.Module):
    """
    Fused Mobile Inverted Bottleneck Convolution.
    
    Replaces the Expand + Depthwise pattern with a single
    regular convolution for better hardware utilization.
    
    Structure:
    Input -> Fused Expand+Spatial (3x3) -> SE -> Project (1x1) -> Output
         \-----------------------------------------------/
                        (residual if in_ch == out_ch)
    """
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: int = 3,
        stride: int = 1,
        expand_ratio: int = 4,  # Smaller than standard MBConv
        se_ratio: float = 0.25
    ):
        super().__init__()
        self.stride = stride
        self.use_residual = (stride == 1) and (in_channels == out_channels)
        
        expanded_channels = in_channels * expand_ratio
        
        layers = []
        
        # 1. Fused expansion + spatial convolution
        layers.extend([
            nn.Conv2d(
                in_channels, expanded_channels, kernel_size,
                stride=stride, padding=kernel_size // 2, bias=False
            ),
            nn.BatchNorm2d(expanded_channels),
            nn.SiLU(inplace=True)
        ])
        
        # 2. Squeeze-and-Excitation
        reduced_dim = max(1, int(in_channels * se_ratio))
        layers.append(SqueezeExcitation(expanded_channels, reduced_dim))
        
        # 3. Projection (linear, no activation)
        layers.extend([
            nn.Conv2d(expanded_channels, out_channels, 1, bias=False),
            nn.BatchNorm2d(out_channels)
        ])
        
        self.block = nn.Sequential(*layers)
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        out = self.block(x)
        if self.use_residual:
            out = out + x
        return out
 
 
class ProgressiveResizing:
    """
    Progressive resizing scheduler for training.
    
    Starts with small images and progressively increases
    resolution throughout training.
    """
    def __init__(
        self,
        initial_size: int = 128,
        final_size: int = 380,
        total_epochs: int = 350,
        ramp_epochs: int = 270  # When to reach final size
    ):
        self.initial_size = initial_size
        self.final_size = final_size
        self.total_epochs = total_epochs
        self.ramp_epochs = ramp_epochs
    
    def get_size(self, epoch: int) -> int:
        """Get image size for current epoch."""
        if epoch >= self.ramp_epochs:
            return self.final_size
        
        # Linear interpolation
        progress = epoch / self.ramp_epochs
        size = self.initial_size + (self.final_size - self.initial_size) * progress
        
        # Round to multiple of 32 for efficient processing
        return int(size // 32 * 32)
    
    def get_regularization_strength(self, epoch: int) -> float:
        """
        Scale regularization with image size.
        
        Smaller images need less regularization (less chance of overfitting);
        larger images need more.
        """
        size = self.get_size(epoch)
        # Linearly scale dropout, mixup, etc.
        scale = (size - self.initial_size) / (self.final_size - self.initial_size)
        return scale

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
114

import torch
import torchvision.transforms as T
from PIL import Image
 
# EfficientNet-specific preprocessing
# Note: Resolution must match model variant!
EFFICIENTNET_CONFIGS = {
    'b0': {'size': 224, 'crop': 224},
    'b1': {'size': 240, 'crop': 240},
    'b2': {'size': 260, 'crop': 260},
    'b3': {'size': 300, 'crop': 300},
    'b4': {'size': 380, 'crop': 380},
    'b5': {'size': 456, 'crop': 456},
    'b6': {'size': 528, 'crop': 528},
    'b7': {'size': 600, 'crop': 600},
}
 
def get_efficientnet_transforms(variant: str = 'b0', training: bool = False):
    """Get appropriate transforms for EfficientNet variant."""
    config = EFFICIENTNET_CONFIGS[variant]
    
    # ImageNet normalization (used during training)
    normalize = T.Normalize(
        mean=[0.485, 0.456, 0.406],
        std=[0.229, 0.224, 0.225]
    )
    
    if training:
        # Training: random resize crop + augmentation
        return T.Compose([
            T.RandomResizedCrop(config['crop']),
            T.RandomHorizontalFlip(),
            T.AutoAugment(T.AutoAugmentPolicy.IMAGENET),
            T.ToTensor(),
            normalize,
        ])
    else:
        # Inference: deterministic resize + center crop
        # Resize to slightly larger, then crop to match training
        resize_size = int(config['size'] / 0.875)  # ~= size / (224/256)
        return T.Compose([
            T.Resize(resize_size, interpolation=T.InterpolationMode.BICUBIC),
            T.CenterCrop(config['crop']),
            T.ToTensor(),
            normalize,
        ])
 
 
class EfficientNetPredictor:
    """Optimized EfficientNet inference wrapper."""
    
    def __init__(self, model, variant: str = 'b0', device: str = 'cuda'):
        self.device = torch.device(device)
        self.model = model.to(self.device).eval()
        self.transforms = get_efficientnet_transforms(variant, training=False)
        
        # Enable inference optimizations
        if device == 'cuda':
            self.model = self.model.half()  # FP16 inference
            torch.backends.cudnn.benchmark = True
    
    @torch.inference_mode()
    def predict(self, image: Image.Image) -> torch.Tensor:
        """Run inference on a single image."""
        x = self.transforms(image).unsqueeze(0)  # Add batch dimension
        x = x.to(self.device)
        
        if self.device.type == 'cuda':
            x = x.half()
        
        logits = self.model(x)
        probabilities = torch.softmax(logits, dim=1)
        return probabilities
    
    @torch.inference_mode()
    def predict_batch(self, images: list) -> torch.Tensor:
        """Run inference on a batch of images."""
        tensors = [self.transforms(img) for img in images]
        batch = torch.stack(tensors).to(self.device)
        
        if self.device.type == 'cuda':
            batch = batch.half()
        
        logits = self.model(batch)
        probabilities = torch.softmax(logits, dim=1)
        return probabilities
 
 
# TensorRT optimization example (pseudocode)
def optimize_for_tensorrt(model, variant: str = 'b0'):
    """
    Convert EfficientNet to TensorRT for maximum inference speed.
    
    This can provide 2-4x speedup on NVIDIA GPUs.
    """
    import torch_tensorrt
    
    config = EFFICIENTNET_CONFIGS[variant]
    
    # Define input specification
    input_spec = torch_tensorrt.Input(
        shape=(1, 3, config['crop'], config['crop']),
        dtype=torch.half
    )
    
    # Compile with TensorRT
    optimized = torch_tensorrt.compile(
        model,
        inputs=[input_spec],
        enabled_precisions={torch.half},
        workspace_size=1 << 30,  # 1 GB workspace
    )
    
    return optimized