import torch
import torch.nn as nn
import torch.nn.functional as Fclass 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_dimself.hidden_dim = hidden_dimdef 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
