PCFN

import torch
import torch.nn as nn
import torch.nn.functional as F

class PCFN(nn.Module):
    '''
    使用带有GELU的激活函数的1*1卷积对扩展的隐藏空间进行跨信道交互。 然后将隐藏特征分割成两块 对其中一块使用3*3卷积核GELU激活函数 编码局部上下文
    将处理后的结果和另一块合并
    '''
    def __init__(self, dim, growth_rate=2.0, p_rate=0.25):
        super().__init__()
        hidden_dim = int(dim * growth_rate)
        p_dim = int(hidden_dim * p_rate)
        self.conv_0 = nn.Conv2d(dim, hidden_dim, 1, 1, 0)
        self.conv_1 = nn.Conv2d(p_dim, p_dim, 3, 1, 1)

        self.act = nn.GELU()
        self.conv_2 = nn.Conv2d(hidden_dim, dim, 1, 1, 0)

        self.p_dim = p_dim
        self.hidden_dim = hidden_dim

    def forward(self, x):
        if self.training:
            '''
            split 和 cat操作都会开辟新的内存
            '''
            x = self.act(self.conv_0(x))
            x1, x2 = torch.split(x, [self.p_dim, self.hidden_dim - self.p_dim], dim=1)
            x1 = self.act(self.conv_1(x1))
            x = self.conv_2(torch.cat([x1, x2], dim=1))
        else:
            '''
            所有的都是原地操作 更节省内存
            '''
            x = self.act(self.conv_0(x))
            x[:, :self.p_dim, :, :] = self.act(self.conv_1(x[:, :self.p_dim, :, :]))
            x = self.conv_2(x)
        return x