classification-main/model/layer.py

import torch
import sys
sys.path.append('/home/yhsun/classification-main/')

class cbs(torch.nn.Module):
    def __init__(self, in_, out_, kernel_size, stride):
        super().__init__()
        self.conv = torch.nn.Conv2d(in_, out_, kernel_size=kernel_size, stride=stride, padding=(kernel_size - 1) // 2,
                                    bias=False)
        self.bn = torch.nn.BatchNorm2d(out_, eps=0.001, momentum=0.03)
        self.silu = torch.nn.SiLU(inplace=True)

    def forward(self, x):
        x = self.conv(x)
        x = self.bn(x)
        x = self.silu(x)
        return x


class concat(torch.nn.Module):
    def __init__(self, dim=1):
        super().__init__()
        self.concat = torch.cat
        self.dim = dim

    def forward(self, x):
        x = self.concat(x, dim=self.dim)
        return x


class elan(torch.nn.Module):  # in_->out_，len->len
    def __init__(self, in_, out_, n, config=None):
        super().__init__()
        if not config:  # 正常版本
            self.cbs0 = cbs(in_, out_ // 4, kernel_size=1, stride=1)
            self.cbs1 = cbs(in_, out_ // 4, kernel_size=1, stride=1)
            self.sequential2 = torch.nn.Sequential(
                *(cbs(out_ // 4, out_ // 4, kernel_size=3, stride=1) for _ in range(n)))
            self.sequential3 = torch.nn.Sequential(
                *(cbs(out_ // 4, out_ // 4, kernel_size=3, stride=1) for _ in range(n)))
            self.concat4 = concat()
            self.cbs5 = cbs(out_, out_, kernel_size=1, stride=1)
        else:  # 剪枝版本。len(config) = 3 + 2 * n
            self.cbs0 = cbs(in_, config[0], kernel_size=1, stride=1)
            self.cbs1 = cbs(in_, config[1], kernel_size=1, stride=1)
            self.sequential2 = torch.nn.Sequential(
                *(cbs(config[1 + _], config[2 + _], kernel_size=3, stride=1) for _ in range(n)))
            self.sequential3 = torch.nn.Sequential(
                *(cbs(config[1 + n + _], config[2 + n + _], kernel_size=3, stride=1) for _ in range(n)))
            self.concat4 = concat()
            self.cbs5 = cbs(config[0] + config[1] + config[1 + n] + config[1 + 2 * n], config[2 + 2 * n],
                            kernel_size=1, stride=1)

    def forward(self, x):
        x0 = self.cbs0(x)
        x1 = self.cbs1(x)
        x2 = self.sequential2(x1)
        x3 = self.sequential3(x2)
        x = self.concat4([x0, x1, x2, x3])
        x = self.cbs5(x)
        return x


class mp(torch.nn.Module):  # in_->out_，len->len//2
    def __init__(self, in_, out_, config=None):
        super().__init__()
        if not config:  # 正常版本
            self.maxpool0 = torch.nn.MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1)
            self.cbs1 = cbs(in_, out_ // 2, 1, 1)
            self.cbs2 = cbs(in_, out_ // 2, 1, 1)
            self.cbs3 = cbs(out_ // 2, out_ // 2, 3, 2)
            self.concat4 = concat(dim=1)
        else:  # 剪枝版本。len(config) = 3
            self.maxpool0 = torch.nn.MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1)
            self.cbs1 = cbs(in_, config[0], 1, 1)
            self.cbs2 = cbs(in_, config[1], 1, 1)
            self.cbs3 = cbs(config[1], config[2], 3, 2)
            self.concat4 = concat(dim=1)

    def forward(self, x):
        x0 = self.maxpool0(x)
        x0 = self.cbs1(x0)
        x1 = self.cbs2(x)
        x1 = self.cbs3(x1)
        x = self.concat4([x0, x1])
        return x


class sppcspc(torch.nn.Module):  # in_->out_，len->len
    def __init__(self, in_, out_, config=None):
        super().__init__()
        if not config:  # 正常版本
            self.cbs0 = cbs(in_, in_ // 2, kernel_size=1, stride=1)
            self.cbs1 = cbs(in_, in_ // 2, kernel_size=1, stride=1)
            self.cbs2 = cbs(in_ // 2, in_ // 2, kernel_size=3, stride=1)
            self.cbs3 = cbs(in_ // 2, in_ // 2, kernel_size=1, stride=1)
            self.MaxPool2d4 = torch.nn.MaxPool2d(kernel_size=5, stride=1, padding=2, dilation=1)
            self.MaxPool2d5 = torch.nn.MaxPool2d(kernel_size=9, stride=1, padding=4, dilation=1)
            self.MaxPool2d6 = torch.nn.MaxPool2d(kernel_size=13, stride=1, padding=6, dilation=1)
            self.concat7 = concat(dim=1)
            self.cbs8 = cbs(2 * in_, in_ // 2, kernel_size=1, stride=1)
            self.cbs9 = cbs(in_ // 2, in_ // 2, kernel_size=3, stride=1)
            self.concat10 = concat(dim=1)
            self.cbs11 = cbs(in_, out_, kernel_size=1, stride=1)
        else:  # 剪枝版本。len(config) = 7
            self.cbs0 = cbs(in_, config[0], kernel_size=1, stride=1)
            self.cbs1 = cbs(in_, config[1], kernel_size=1, stride=1)
            self.cbs2 = cbs(config[1], config[2], kernel_size=3, stride=1)
            self.cbs3 = cbs(config[2], config[3], kernel_size=1, stride=1)
            self.MaxPool2d4 = torch.nn.MaxPool2d(kernel_size=5, stride=1, padding=2, dilation=1)
            self.MaxPool2d5 = torch.nn.MaxPool2d(kernel_size=9, stride=1, padding=4, dilation=1)
            self.MaxPool2d6 = torch.nn.MaxPool2d(kernel_size=13, stride=1, padding=6, dilation=1)
            self.concat7 = concat(dim=1)
            self.cbs8 = cbs(4 * config[3], config[4], kernel_size=1, stride=1)
            self.cbs9 = cbs(config[4], config[5], kernel_size=3, stride=1)
            self.concat10 = concat(dim=1)
            self.cbs11 = cbs(config[0] + config[5], config[6], kernel_size=1, stride=1)

    def forward(self, x):
        x0 = self.cbs0(x)
        x1 = self.cbs1(x)
        x1 = self.cbs2(x1)
        x1 = self.cbs3(x1)
        x4 = self.MaxPool2d4(x1)
        x5 = self.MaxPool2d5(x1)
        x6 = self.MaxPool2d6(x1)
        x = self.concat7([x1, x4, x5, x6])
        x = self.cbs8(x)
        x = self.cbs9(x)
        x = self.concat10([x, x0])
        x = self.cbs11(x)
        return x


class linear_head(torch.nn.Module):
    def __init__(self, in_, out_):
        super().__init__()
        self.avgpool0 = torch.nn.AdaptiveAvgPool2d(1)
        self.flatten1 = torch.nn.Flatten()
        self.Dropout2 = torch.nn.Dropout(0.2)
        self.linear3 = torch.nn.Linear(in_, in_ // 2)
        self.silu4 = torch.nn.SiLU()
        self.Dropout5 = torch.nn.Dropout(0.2)
        self.linear6 = torch.nn.Linear(in_ // 2, out_)

    def forward(self, x):
        x = self.avgpool0(x)
        x = self.flatten1(x)
        x = self.Dropout2(x)
        x = self.linear3(x)
        x = self.silu4(x)
        x = self.Dropout5(x)
        x = self.linear6(x)
        return x


class image_deal(torch.nn.Module):  # 归一化
    def __init__(self):
        super().__init__()

    def forward(self, x):
        x = x / 255
        x = x.permute(0, 3, 1, 2)
        return x


class deploy(torch.nn.Module):
    def __init__(self, model, normalization):
        super().__init__()
        self.image_deal = image_deal()
        self.model = model
        if normalization == 'softmax':
            self.normalization = torch.nn.Softmax(dim=1)
        else:
            self.normalization = torch.nn.Sigmoid()

    def forward(self, x):
        x = self.image_deal(x)
        x = self.model(x)
        x = self.normalization(x)
        return x