# -*- coding: utf-8 -*- # @Author: Weisen Pan import torch import torch.nn as nn # Try to import the method to load model weights from a URL, with a fallback in case of ImportError try: from torch.hub import load_state_dict_from_url except ImportError: from torch.utils.model_zoo import load_url as load_state_dict_from_url # List of available ResNet architectures __all__ = ['resnet_model_18', 'resnet_model_34', 'resnet_model_50', 'resnet_model_101', 'resnet_model_152', 'resnet_model_200', 'resnet110', 'resnet164', 'resnext29_8x64d', 'resnext29_16x64d', 'resnext50_32x4d', 'resnext101_32x4d', 'resnext101_32x8d', 'resnext101_64x4d', 'wide_resnet_model_50_2', 'wide_resnet_model_50_3', 'wide_resnet_model_101_2', 'wide_resnet16_8', 'wide_resnet52_8', 'wide_resnet16_12', 'wide_resnet28_10', 'wide_resnet40_10'] # Pre-trained model URLs for various ResNet variants model_urls = { 'resnet_model_18': 'https://download.pytorch.org/models/resnet_model_18-5c106cde.pth', 'resnet_model_34': 'https://download.pytorch.org/models/resnet_model_34-333f7ec4.pth', 'resnet_model_50': 'https://download.pytorch.org/models/resnet_model_50-19c8e357.pth', 'resnet_model_101': 'https://download.pytorch.org/models/resnet_model_101-5d3b4d8f.pth', 'resnet_model_152': 'https://download.pytorch.org/models/resnet_model_152-b121ed2d.pth', 'resnext50_32x4d': 'https://download.pytorch.org/models/resnext50_32x4d-7cdf4587.pth', 'resnext101_32x8d': 'https://download.pytorch.org/models/resnext101_32x8d-8ba56ff5.pth', 'wide_resnet_model_50_2': 'https://download.pytorch.org/models/wide_resnet_model_50_2-95faca4d.pth', 'wide_resnet_model_101_2': 'https://download.pytorch.org/models/wide_resnet_model_101_2-32ee1156.pth', } # Function for a 3x3 convolution with padding def apply_3x3_convolution(in_planes, out_planes, stride=1, groups=1, dilation=1): """3x3 convolution with padding""" return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, padding=dilation, groups=groups, bias=False, dilation=dilation) # Function for a 1x1 convolution def apply_1x1_convolution(in_planes, out_planes, stride=1): """1x1 convolution""" return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False) # BasicBlock class for the ResNet architecture class BasicBlock(nn.Module): expansion = 1 # Expansion factor for the output channels def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1, base_width=64, dilation=1, norm_layer=None): super(BasicBlock, self).__init__() # If norm_layer is not provided, use BatchNorm2d as the default if norm_layer is None: norm_layer = nn.BatchNorm2d # Ensure BasicBlock is restricted to specific parameters if groups != 1 or base_width != 64: raise ValueError('BasicBlock is restricted to groups=1 and base_width=64') if dilation > 1: raise NotImplementedError("BasicBlock does not support dilation greater than 1") # Define the layers for the BasicBlock self.conv1 = apply_3x3_convolution(inplanes, planes, stride) # First 3x3 convolution self.bn1 = norm_layer(planes) # First BatchNorm layer self.relu = nn.ReLU(inplace=True) # ReLU activation self.conv2 = apply_3x3_convolution(planes, planes) # Second 3x3 convolution self.bn2 = norm_layer(planes) # Second BatchNorm layer self.downsample = downsample # Optional downsample layer self.stride = stride # Define the forward pass for BasicBlock def forward(self, x): identity = x # Save the input for the skip connection out = self.conv1(x) # First convolution out = self.bn1(out) # BatchNorm after first convolution out = self.relu(out) # ReLU activation out = self.conv2(out) # Second convolution out = self.bn2(out) # BatchNorm after second convolution # Apply downsample if defined if self.downsample is not None: identity = self.downsample(x) out += identity # Add the skip connection out = self.relu(out) # Apply ReLU activation again return out # Bottleneck class for the ResNet architecture, a more complex block used in deeper ResNet models class Bottleneck(nn.Module): expansion = 4 # Expansion factor for the output channels def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1, base_width=64, dilation=1, norm_layer=None): super(Bottleneck, self).__init__() if norm_layer is None: norm_layer = nn.BatchNorm2d width = int(planes * (base_width / 64.)) * groups # Calculate width based on base width and groups # Define the layers for the Bottleneck block self.conv1 = apply_1x1_convolution(inplanes, width) # 1x1 convolution to reduce the dimensions self.bn1 = norm_layer(width) # BatchNorm after 1x1 convolution self.conv2 = apply_3x3_convolution(width, width, stride, groups, dilation) # 3x3 convolution self.bn2 = norm_layer(width) # BatchNorm after 3x3 convolution self.conv3 = apply_1x1_convolution(width, planes * self.expansion) # 1x1 convolution to expand the dimensions self.bn3 = norm_layer(planes * self.expansion) # BatchNorm after final 1x1 convolution self.relu = nn.ReLU(inplace=True) # ReLU activation self.downsample = downsample # Optional downsample layer self.stride = stride # Define the forward pass for Bottleneck def forward(self, x): identity = x # Save the input for the skip connection out = self.conv1(x) # First convolution out = self.bn1(out) # BatchNorm after first convolution out = self.relu(out) # ReLU activation out = self.conv2(out) # Second convolution out = self.bn2(out) # BatchNorm after second convolution out = self.relu(out) # ReLU activation out = self.conv3(out) # Third convolution out = self.bn3(out) # BatchNorm after third convolution # Apply downsample if defined if self.downsample is not None: identity = self.downsample(x) out += identity # Add the skip connection out = self.relu(out) # Apply ReLU activation again return out # Main ResNet class, a customizable deep learning model architecture class ResNet(nn.Module): def __init__(self, arch, block, layers, num_classes=1000, zero_init_residual=True, groups=1, width_per_group=64, replace_stride_with_dilation=None, norm_layer=None, dataset='cifar10', split_factor=1, output_stride=8, dropout_p=None): super(ResNet, self).__init__() if norm_layer is None: norm_layer = nn.BatchNorm2d # Default normalization layer self._norm_layer = norm_layer self.groups = groups # Number of groups in convolutions self.inplanes = 16 if dataset in ['cifar10', 'cifar100'] else 64 # Adjust initial planes for CIFAR # First layer: a combination of convolution, normalization, and ReLU self.layer0 = nn.Sequential( nn.Conv2d(3, self.inplanes, kernel_size=3, stride=1, padding=1, bias=False), norm_layer(self.inplanes), nn.ReLU(inplace=True), ) # Subsequent ResNet layers using the _create_model_layer method self.layer1 = self._create_model_layer(block, 16, layers[0]) self.layer2 = self._create_model_layer(block, 32, layers[1], stride=2) self.layer3 = self._create_model_layer(block, 64, layers[2], stride=2) self.avgpool = nn.AdaptiveAvgPool2d((1, 1)) # Global average pooling self.fc = nn.Linear(64 * block.expansion, num_classes) # Fully connected layer for classification # Initialization for model weights for m in self.modules(): if isinstance(m, nn.Conv2d): nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)): nn.init.constant_(m.weight, 1) nn.init.constant_(m.bias, 0) elif isinstance(m, nn.Linear): nn.init.normal_(m.weight, 0, 1e-3) # Zero-initialize the last BatchNorm in residual connections if required if zero_init_residual: for m in self.modules(): if isinstance(m, Bottleneck): nn.init.constant_(m.bn3.weight, 0) elif isinstance(m, BasicBlock): nn.init.constant_(m.bn2.weight, 0) # Helper function to create layers in ResNet def _create_model_layer(self, block, planes, blocks, stride=1, dilate=False): norm_layer = self._norm_layer # Set normalization layer downsample = None # If the stride is not 1 or input/output planes do not match, create a downsample layer if stride != 1 or self.inplanes != planes * block.expansion: downsample = nn.Sequential( apply_1x1_convolution(self.inplanes, planes * block.expansion, stride), norm_layer(planes * block.expansion), ) layers = [block(self.inplanes, planes, stride, downsample)] # Create the first block with downsampling self.inplanes = planes * block.expansion # Update inplanes for next blocks for _ in range(1, blocks): layers.append(block(self.inplanes, planes)) # Add subsequent blocks without downsampling return nn.Sequential(*layers) # Forward pass through the ResNet architecture def forward(self, x): x = self.layer0(x) # Pass input through the first layer x = self.layer1(x) # First ResNet layer x = self.layer2(x) # Second ResNet layer x = self.layer3(x) # Third ResNet layer x = self.avgpool(x) # Global average pooling x = torch.flatten(x, 1) # Flatten the output for the fully connected layer x = self.fc(x) # Pass through the fully connected layer return x # Helper function to instantiate ResNet with pretrained weights if available def _resnet(arch, block, layers, models_pretrained, progress, **kwargs): model = ResNet(arch, block, layers, **kwargs) # Create a ResNet model if models_pretrained: # Load pretrained weights if requested state_dict = load_state_dict_from_url(model_urls[arch], progress=progress) model.load_state_dict(state_dict) return model # Functions to create specific ResNet variants def resnet_model_18(models_pretrained=False, progress=True, **kwargs): return _resnet('resnet_model_18', BasicBlock, [2, 2, 2, 2], models_pretrained, progress, **kwargs) def resnet_model_34(models_pretrained=False, progress=True, **kwargs): return _resnet('resnet_model_34', BasicBlock, [3, 4, 6, 3], models_pretrained, progress, **kwargs) def resnet_model_50(models_pretrained=False, progress=True, **kwargs): return _resnet('resnet_model_50', Bottleneck, [3, 4, 6, 3], models_pretrained, progress, **kwargs) def resnet_model_101(models_pretrained=False, progress=True, **kwargs): return _resnet('resnet_model_101', Bottleneck, [3, 4, 23, 3], models_pretrained, progress, **kwargs) def resnet_model_152(models_pretrained=False, progress=True, **kwargs): return _resnet('resnet_model_152', Bottleneck, [3, 8, 36, 3], models_pretrained, progress, **kwargs) def resnet_model_200(models_pretrained=False, progress=True, **kwargs): return _resnet('resnet_model_200', Bottleneck, [3, 24, 36, 3], models_pretrained, progress, **kwargs)