updated and checked CNN architectures still works with latest pytorch

ML/Pytorch/CNN_architectures/pytorch_efficientnet.py

@@ -1,3 +1,12 @@
+"""
+An implementation of EfficientNet CNN architecture.
+
+Programmed by Aladdin Persson <aladdin.persson at hotmail dot com>
+*    2021-02-05 Initial coding
+*    2022-12-20 Update comments, code revision, checked still works with latest PyTorch version
+"""
+
+
 import torch
 import torch.nn as nn
 from math import ceil
@@ -25,9 +34,10 @@ phi_values = {
     "b7": (6, 600, 0.5),
 }
 
+
 class CNNBlock(nn.Module):
     def __init__(
-            self, in_channels, out_channels, kernel_size, stride, padding, groups=1
+        self, in_channels, out_channels, kernel_size, stride, padding, groups=1
     ):
         super(CNNBlock, self).__init__()
         self.cnn = nn.Conv2d(
@@ -40,16 +50,17 @@ class CNNBlock(nn.Module):
             bias=False,
         )
         self.bn = nn.BatchNorm2d(out_channels)
-        self.silu = nn.SiLU() # SiLU <-> Swish
+        self.silu = nn.SiLU()  # SiLU <-> Swish
 
     def forward(self, x):
         return self.silu(self.bn(self.cnn(x)))
 
+
 class SqueezeExcitation(nn.Module):
     def __init__(self, in_channels, reduced_dim):
         super(SqueezeExcitation, self).__init__()
         self.se = nn.Sequential(
-            nn.AdaptiveAvgPool2d(1), # C x H x W -> C x 1 x 1
+            nn.AdaptiveAvgPool2d(1),  # C x H x W -> C x 1 x 1
             nn.Conv2d(in_channels, reduced_dim, 1),
             nn.SiLU(),
             nn.Conv2d(reduced_dim, in_channels, 1),
@@ -59,17 +70,18 @@ class SqueezeExcitation(nn.Module):
     def forward(self, x):
         return x * self.se(x)
 
+
 class InvertedResidualBlock(nn.Module):
     def __init__(
-            self,
-            in_channels,
-            out_channels,
-            kernel_size,
-            stride,
-            padding,
-            expand_ratio,
-            reduction=4, # squeeze excitation
-            survival_prob=0.8, # for stochastic depth
+        self,
+        in_channels,
+        out_channels,
+        kernel_size,
+        stride,
+        padding,
+        expand_ratio,
+        reduction=4,  # squeeze excitation
+        survival_prob=0.8,  # for stochastic depth
     ):
         super(InvertedResidualBlock, self).__init__()
         self.survival_prob = 0.8
@@ -80,12 +92,21 @@ class InvertedResidualBlock(nn.Module):
 
         if self.expand:
             self.expand_conv = CNNBlock(
-                in_channels, hidden_dim, kernel_size=3, stride=1, padding=1,
+                in_channels,
+                hidden_dim,
+                kernel_size=3,
+                stride=1,
+                padding=1,
             )
 
         self.conv = nn.Sequential(
             CNNBlock(
-                hidden_dim, hidden_dim, kernel_size, stride, padding, groups=hidden_dim,
+                hidden_dim,
+                hidden_dim,
+                kernel_size,
+                stride,
+                padding,
+                groups=hidden_dim,
             ),
             SqueezeExcitation(hidden_dim, reduced_dim),
             nn.Conv2d(hidden_dim, out_channels, 1, bias=False),
@@ -96,7 +117,9 @@ class InvertedResidualBlock(nn.Module):
         if not self.training:
             return x
 
-        binary_tensor = torch.rand(x.shape[0], 1, 1, 1, device=x.device) < self.survival_prob
+        binary_tensor = (
+            torch.rand(x.shape[0], 1, 1, 1, device=x.device) < self.survival_prob
+        )
         return torch.div(x, self.survival_prob) * binary_tensor
 
     def forward(self, inputs):
@@ -122,8 +145,8 @@ class EfficientNet(nn.Module):
 
     def calculate_factors(self, version, alpha=1.2, beta=1.1):
         phi, res, drop_rate = phi_values[version]
-        depth_factor = alpha ** phi
-        width_factor = beta ** phi
+        depth_factor = alpha**phi
+        width_factor = beta**phi
         return width_factor, depth_factor, drop_rate
 
     def create_features(self, width_factor, depth_factor, last_channels):
@@ -132,7 +155,7 @@ class EfficientNet(nn.Module):
         in_channels = channels
 
         for expand_ratio, channels, repeats, stride, kernel_size in base_model:
-            out_channels = 4*ceil(int(channels*width_factor) / 4)
+            out_channels = 4 * ceil(int(channels * width_factor) / 4)
             layers_repeats = ceil(repeats * depth_factor)
 
             for layer in range(layers_repeats):
@@ -141,9 +164,9 @@ class EfficientNet(nn.Module):
                         in_channels,
                         out_channels,
                         expand_ratio=expand_ratio,
-                        stride = stride if layer == 0 else 1,
+                        stride=stride if layer == 0 else 1,
                         kernel_size=kernel_size,
-                        padding=kernel_size//2, # if k=1:pad=0, k=3:pad=1, k=5:pad=2
+                        padding=kernel_size // 2,  # if k=1:pad=0, k=3:pad=1, k=5:pad=2
                     )
                 )
                 in_channels = out_channels
@@ -170,6 +193,8 @@ def test():
         num_classes=num_classes,
     ).to(device)
 
-    print(model(x).shape) # (num_examples, num_classes)
+    print(model(x).shape)  # (num_examples, num_classes)
 
-test()
+
+if __name__ == "__main__":
+    test()