import torch
import torch.nn as nn
from torchsummary import summary

class Residual(nn.Module):
def init(self,in_channel,out_channel,use_1conv = False,strides=1):
super(Residual,self).init()
self.Relu = nn.ReLU()
#################conv1和CONV2的 stride必须是相同的，才能保证相同的像素宽度和高度
self.conv1 = nn.Conv2d(in_channels=in_channel,out_channels=out_channel,kernel_size=3,padding=1,stride=strides)
self.conv2= nn.Conv2d(in_channels=out_channel,out_channels=out_channel,kernel_size=3,padding=1)
self.bn1 = nn.BatchNorm2d(out_channel)
self.bn2 = nn.BatchNorm2d(out_channel)
if use_1conv:
self.conv3 = nn.Conv2d(in_channels=in_channel,out_channels=out_channel,kernel_size=1,stride=strides)
else:
self.conv3 = None

def forward(self,x):y= self.Relu(self.bn1(self.conv1(x)))y =self.bn2(self.conv2(y))if self.conv3:x = self.conv3(x)y=self.Relu(y+x)return y

class ResNet18(nn.Module):
def init(self,Residual):
super(ResNet18,self).init()
self.b1 = nn.Sequential(
nn.Conv2d(in_channels=1,out_channels=64,kernel_size=7,stride=2,padding=3),
nn.ReLU(),
nn.BatchNorm2d(64),
nn.MaxPool2d(kernel_size=3,stride=2,padding=1)
)

    self.b2 = nn.Sequential(Residual(64,64,use_1conv = False,strides=1),Residual(64,64,use_1conv = False,strides=1),)self.b3 =nn.Sequential(Residual(64, 128, use_1conv=True, strides=2),Residual(128, 128, use_1conv=False, strides=1),)self.b4 = nn.Sequential(Residual(128, 256, use_1conv=True, strides=2),Residual(256, 256, use_1conv=False, strides=1),)self.b5 = nn.Sequential(Residual(256, 512, use_1conv=True, strides=1),Residual(512, 512, use_1conv=False, strides=1),)self.b6 =nn.Sequential(nn.AdaptiveAvgPool2d((1,1)),nn.Flatten(),nn.Linear(512,10))
def forward(self,x):x = self.b1(x)x = self.b2(x)x = self.b3(x)x = self.b4(x)x = self.b5(x)x = self.b6(x)return x

if name==“main”:
device = torch.device(“cuda” if torch.cuda.is_available() else “cpu”)
model1 = ResNet18(Residual).to(device)
print(summary(model1,(1,224,224)))