| 84 | return x |
| 85 | |
| 86 | class EfficientNetBlock(nn.Module): |
| 87 | def __init__(self, in_channels, out_channels, kernel_size, stride, expand_ratio, se_ratio=0.25): |
| 88 | super().__init__() |
| 89 | self.stride = stride |
| 90 | self.use_residual = (stride == 1 and in_channels == out_channels) |
| 91 | |
| 92 | hidden_dim = in_channels * expand_ratio |
| 93 | self.use_expansion = expand_ratio != 1 |
| 94 | |
| 95 | if self.use_expansion: |
| 96 | self.expand_conv = nn.Sequential( |
| 97 | nn.Conv2d(in_channels, hidden_dim, 1, bias=False), |
| 98 | nn.BatchNorm2d(hidden_dim), |
| 99 | nn.SiLU(inplace=True) |
| 100 | ) |
| 101 | |
| 102 | self.depthwise_conv = nn.Sequential( |
| 103 | nn.Conv2d(hidden_dim, hidden_dim, kernel_size, stride, kernel_size // 2, groups=hidden_dim, bias=False), |
| 104 | nn.BatchNorm2d(hidden_dim), |
| 105 | nn.SiLU(inplace=True) |
| 106 | ) |
| 107 | |
| 108 | se_channels = max(1, int(in_channels * se_ratio)) |
| 109 | self.se = nn.Sequential( |
| 110 | nn.AdaptiveAvgPool2d(1), |
| 111 | nn.Conv2d(hidden_dim, se_channels, 1), |
| 112 | nn.SiLU(inplace=True), |
| 113 | nn.Conv2d(se_channels, hidden_dim, 1), |
| 114 | nn.Sigmoid() |
| 115 | ) |
| 116 | |
| 117 | self.project_conv = nn.Sequential( |
| 118 | nn.Conv2d(hidden_dim, out_channels, 1, bias=False), |
| 119 | nn.BatchNorm2d(out_channels) |
| 120 | ) |
| 121 | |
| 122 | def forward(self, x): |
| 123 | identity = x |
| 124 | |
| 125 | if self.use_expansion: |
| 126 | x = self.expand_conv(x) |
| 127 | |
| 128 | x = self.depthwise_conv(x) |
| 129 | |
| 130 | se_weight = self.se(x) |
| 131 | x = x * se_weight |
| 132 | |
| 133 | x = self.project_conv(x) |
| 134 | |
| 135 | if self.use_residual: |
| 136 | x = x + identity |
| 137 | |
| 138 | return x |
| 139 | |
| 140 | class UNetBlock(nn.Module): |
| 141 | def __init__(self, in_channels, out_channels, down=True): |