Feature fusion block.
| 192 | |
| 193 | |
| 194 | class FeatureFusionBlock(nn.Module): |
| 195 | """Feature fusion block. |
| 196 | """ |
| 197 | |
| 198 | def __init__(self, features): |
| 199 | """Init. |
| 200 | |
| 201 | Args: |
| 202 | features (int): number of features |
| 203 | """ |
| 204 | super(FeatureFusionBlock, self).__init__() |
| 205 | |
| 206 | self.resConfUnit1 = ResidualConvUnit(features) |
| 207 | self.resConfUnit2 = ResidualConvUnit(features) |
| 208 | |
| 209 | def forward(self, *xs): |
| 210 | """Forward pass. |
| 211 | |
| 212 | Returns: |
| 213 | tensor: output |
| 214 | """ |
| 215 | output = xs[0] |
| 216 | |
| 217 | if len(xs) == 2: |
| 218 | output += self.resConfUnit1(xs[1]) |
| 219 | |
| 220 | output = self.resConfUnit2(output) |
| 221 | |
| 222 | output = nn.functional.interpolate( |
| 223 | output, scale_factor=2, mode="bilinear", align_corners=True |
| 224 | ) |
| 225 | |
| 226 | return output |
| 227 | |
| 228 | |
| 229 |