# 总的残差网络模型
class ResNet(nn.Module):
    def __init__(self,block,block_num,num_classes=1000,include_top=True):
        super(ResNet,self).__init__()
        self.include_top=include_top
        self.in_channel=64
        # 第一个卷积层，使用7*7的卷积核，步长为2，使数据维度减半
        self.conv1=nn.Conv2d(in_channels=3,out_channels=self.in_channel,
                             kernel_size=7,stride=2,padding=3,bias=False)
        # 将卷积后的数据进行标准化
        self.bn1=nn.BatchNorm2d(self.in_channel)
        self.relu=nn.ReLU(inplace=True)
        # 最大池化层，采用卷积核为3，同样也会是维度减半
        self.maxpool=nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
        '''分别对应4个残差块，这里传参有个注意问题，
        除了layer1，其余传的都是stride=2，因为对应之前说的任何网络的2，3，4，层的卷积块的
        第一个层都要进行数据降维，而layer不用，只需要将卷积核数升高，
        而数据维度不用变
        64,128,256,512就对应每个残差块的最基层的卷积核数，
        50、101层网络每层最终输出是它的4倍'''
        self.layer1=self._make_layer(block,64,block_num[0])
        self.layer2=self._make_layer(block,128,block_num[1],stride=2)
        self.layer3=self._make_layer(block,256,block_num[2],stride=2)
        self.layer4=self._make_layer(block,512,block_num[3],stride=2)
        # 将所有卷积核中的数据进行池化(1,1)之后进行全连接
        if self.include_top:
            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')
    # 创建每一个残差块网络，第一个参数对应使用BisicBlock或者是Bottleneck，
    # 第二个参数是每个残差网络最基层的卷积核层数，第三个是每个残差块有几个基本网络
    def _make_layer(self,block,channel,block_num,stride=1):
        downsample=None
        '''对应那条捷径，50、101的第一层的第一个卷积层需要捷径，
        以及任何网络的第2、3、4个残差网络的第一层需要捷径，
        且50、101第一层的捷径不需要降维，只需要将卷积核的个数升高，
        其它的都需要进行降维且增加卷积核'''
        if stride!=1 or self.in_channel!=channel*block.expansion:
            downsample=nn.Sequential(
                nn.Conv2d(self.in_channel,channel*block.expansion,
                          kernel_size=1,stride=stride,bias=False),
                nn.BatchNorm2d(channel*block.expansion))
        layers=[]
        # 添加第一个卷积层
        layers.append(block(self.in_channel,channel,stride=stride,downsample=downsample))
        self.in_channel=channel*block.expansion
        # 添加之后的层
        for i in range(1,block_num):
            layers.append(block(self.in_channel,channel))
        return nn.Sequential(*layers)
    def forward(self,x):
        out=self.conv1(x)
        out=self.bn1(out)
        out=self.relu(out)
        out=self.maxpool(out)
        out=self.layer1(out)
        out=self.layer2(out)
        out=self.layer3(out)
        out=self.layer4(out)
        if self.include_top:
            out=self.avgpool(out)
            out=torch.flatten(out,1)
            out=self.fc(out)
        return out