text: import torch.nn as nn
  from torch import Tensor
  import torch
  from typing import List, Optional, Type, Union, Callable, Any
  import sys
  sys.path.append('/home/zmyuan/ResNet18-10')
  from ConAM import conam
  
  
  
  def conv3x3(in_planes: int, out_planes: int, stride: int = 1, groups: int = 1, dilation: int = 1) -> nn.Conv2d:
      """3x3 convolution with padding"""
      return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
                       padding=dilation, groups=groups, bias=False, dilation=dilation)
  
  
  def conv1x1(in_planes: int, out_planes: int, stride: int = 1) -> nn.Conv2d:
      """1x1 convolution"""
      return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)
  
  
  class BasicBlock(nn.Module):
      expansion: int = 1
  
      def __init__(
          self,
          inplanes: int,
          planes: int,
          patch_size: int = 0,
          stride: int = 1,
          downsample: Optional[nn.Module] = None,
          groups: int = 1,
          base_width: int = 64,
          dilation: int = 1,
          norm_layer: Optional[Callable[..., nn.Module]] = None
      ) -> None:
          super(BasicBlock, self).__init__()
          if norm_layer is None:
              norm_layer = nn.BatchNorm2d
          if groups != 1 or base_width != 64:
              raise ValueError('BasicBlock only supports groups=1 and base_width=64')
          if dilation > 1:
              raise NotImplementedError("Dilation > 1 not supported in BasicBlock")
          self.patch_size = patch_size
          if patch_size != 0:
              self.attention = conam(patch_size=patch_size, inchannel=inplanes)
          self.conv1 = conv3x3(inplanes, planes, stride)
          self.bn1 = norm_layer(planes)
          self.relu = nn.ReLU(inplace=True)
          self.conv2 = conv3x3(planes, planes)
          self.bn2 = norm_layer(planes)
          self.downsample = downsample
  
      def forward(self, x: Tensor) -> Tensor:
          identity = x
          if self.patch_size != 0:
              x = self.attention(x)
  
          out = self.conv1(x)
          out = self.bn1(out)
          out = self.relu(out)
  
          out = self.conv2(out)
          out = self.bn2(out)
  
          if self.downsample is not None:
              identity = self.downsample(identity)
  
          out += identity
          out = self.relu(out)
  
          return out
  
  class Bottleneck(nn.Module):
  
  
      expansion: int = 4
  
      def __init__(
          self,
          inplanes: int,
          planes: int,
          patch_size: int = 0,
          stride: int = 1,
          downsample: Optional[nn.Module] = None,
          groups: int = 1,
          base_width: int = 64,
          dilation: int = 1,
          norm_layer: Optional[Callable[..., nn.Module]] = None
      ) -> None:
          super(Bottleneck, self).__init__()
          if norm_layer is None:
              norm_layer = nn.BatchNorm2d
          width = int(planes * (base_width / 64.)) * groups
          self.patch_size = patch_size
          if patch_size != 0:
              self.attention = conam(patch_size=patch_size)
          self.conv1 = conv1x1(inplanes, width)
          self.bn1 = norm_layer(width)
          self.conv2 = conv3x3(width, width, stride, groups, dilation)
          self.bn2 = norm_layer(width)
          self.conv3 = conv1x1(width, planes * self.expansion)
          self.bn3 = norm_layer(planes * self.expansion)
          self.relu = nn.ReLU(inplace=True)
          self.downsample = downsample
          self.stride = stride
  
      def forward(self, x: Tensor) -> Tensor:
          identity = x
  
          if self.patch_size != 0:
              x = self.attention(x)
  
          out = self.conv1(x)
          out = self.bn1(out)
          out = self.relu(out)
  
          out = self.conv2(out)
          out = self.bn2(out)
          out = self.relu(out)
  
          out = self.conv3(out)
          out = self.bn3(out)
  
          if self.downsample is not None:
              identity = self.downsample(identity)
  
          out += identity
          out = self.relu(out)
  
          return out
  
  
  class ResNet(nn.Module):
  
      def __init__(
          self,
          block: Type[Union[Bottleneck,BasicBlock]],
          layers: List[int],
          patch_sizes: List[int],
          num_classes: int = 10,
          zero_init_residual: bool = False,
          groups: int = 1,
          width_per_group: int = 64,
          norm_layer: Optional[Callable[..., nn.Module]] = None
      ) -> None:
          super(ResNet, self).__init__()
          if norm_layer is None:
              norm_layer = nn.BatchNorm2d
          self._norm_layer = norm_layer
  
          self.inplanes = 64
          self.dilation = 1
          
          
          self.groups = groups
          self.base_width = width_per_group
          self.conv1 = nn.Conv2d(3, self.inplanes, kernel_size=3, stride=1, padding=1,
                                 bias=False)
          self.bn1 = norm_layer(self.inplanes)
          self.relu = nn.ReLU(inplace=True)
          self.layer1 = self._make_layer(block, 64, layers[0], patch_sizes[0])
          self.layer2 = self._make_layer(block, 128, layers[1], patch_sizes[1], stride=2,)
          self.layer3 = self._make_layer(block, 256, layers[2], patch_sizes[2], stride=2,)
          self.layer4 = self._make_layer(block, 512, layers[3], patch_sizes[3], stride=2,)
          self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
          self.fc = nn.Linear(512 * block.expansion, num_classes)
  
          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)
  
          
          if zero_init_residual:
              for m in self.modules():
                  if isinstance(m, Bottleneck):
                      nn.init.constant_(m.bn3.weight, 0)  
                  else: 
                      raise ValueError('not Bottleneck')
  
      def _make_layer(self, block: Type[Union[Bottleneck, BasicBlock]], planes: int, blocks: int, patch_size: int,
                      stride: int = 1,) -> nn.Sequential:
          norm_layer = self._norm_layer
          downsample = None
          previous_dilation = self.dilation
          if stride != 1 or self.inplanes != planes * block.expansion:
              downsample = nn.Sequential(
                  conv1x1(self.inplanes, planes * block.expansion, stride),
                  norm_layer(planes * block.expansion),
              )
  
          layers = []
          layers.append(block(self.inplanes, planes,  patch_size, stride, downsample, self.groups,
                              self.base_width, previous_dilation, norm_layer))
          self.inplanes = planes * block.expansion
          for _ in range(1, blocks):
              layers.append(block(self.inplanes, planes, groups=self.groups,
                                  base_width=self.base_width, dilation=self.dilation,
                                  norm_layer=norm_layer))
  
          return nn.Sequential(*layers)
  
      def _forward_impl(self, x: Tensor) -> Tensor:
          
          x = self.conv1(x)
          x = self.bn1(x)
          x = self.relu(x)
  
          x = self.layer1(x)
          x = self.layer2(x)
          x = self.layer3(x)
          x = self.layer4(x)
  
          x = self.avgpool(x)
          x = torch.flatten(x, 1)
          x = self.fc(x)
  
          return x
  
      def forward(self, x: Tensor) -> Tensor:
          return self._forward_impl(x)
  
  
  
  
  def resnet18(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> ResNet:
  
      layer = [2, 2, 2, 2]
      patch_size = [8, 8, 4, 0] # the patch_size can be [8, 8, 4, 0], [8, 4, 2, 0], etc..
      return ResNet(BasicBlock, layer, patch_size)