Feature fusion for DPT.
| 125 | |
| 126 | |
| 127 | class FeatureFusionBlock2d(nn.Module): |
| 128 | """Feature fusion for DPT.""" |
| 129 | |
| 130 | # We use the name "deconv" for backward compatibility. However, "deconv" can also |
| 131 | # refer to some other upsampling layer or a no-op. |
| 132 | deconv: nn.Module |
| 133 | |
| 134 | def __init__( |
| 135 | self, |
| 136 | dim_in: int, |
| 137 | dim_out: int | None = None, |
| 138 | upsampling_mode: UpsamplingMode | None = None, |
| 139 | batch_norm: bool = False, |
| 140 | ): |
| 141 | """Initialize feature fusion block. |
| 142 | |
| 143 | Args: |
| 144 | dim_in: Dimensions of input. |
| 145 | dim_out: Dimensions of output. |
| 146 | batch_norm: Whether to use batch normalization in resnet blocks. |
| 147 | upsampling_mode: What mode to use for upsampling. None if no upsampling |
| 148 | is required. |
| 149 | """ |
| 150 | super().__init__() |
| 151 | if dim_out is None: |
| 152 | dim_out = dim_in |
| 153 | self.resnet1 = self._residual_block(dim_in, batch_norm) |
| 154 | self.resnet2 = self._residual_block(dim_in, batch_norm) |
| 155 | |
| 156 | if upsampling_mode is not None: |
| 157 | self.deconv = upsampling_layer(upsampling_mode, scale_factor=2, dim_in=dim_in) |
| 158 | else: |
| 159 | self.deconv = nn.Sequential() |
| 160 | |
| 161 | self.out_conv = nn.Conv2d( |
| 162 | dim_in, |
| 163 | dim_out, |
| 164 | kernel_size=1, |
| 165 | stride=1, |
| 166 | padding=0, |
| 167 | bias=True, |
| 168 | ) |
| 169 | |
| 170 | self.skip_add = nn.quantized.FloatFunctional() |
| 171 | |
| 172 | def forward(self, x0: torch.Tensor, x1: torch.Tensor | None = None) -> torch.Tensor: |
| 173 | """Process and fuse input features.""" |
| 174 | x = x0 |
| 175 | |
| 176 | if x1 is not None: |
| 177 | res = self.resnet1(x1) |
| 178 | x = self.skip_add.add(x, res) |
| 179 | |
| 180 | x = self.resnet2(x) |
| 181 | x = self.deconv(x) |
| 182 | x = self.out_conv(x) |
| 183 | |
| 184 | return x |