A wrapper around a conv layer that behaves like a BaseBackbone.
| 91 | |
| 92 | |
| 93 | class SkipConvBackbone(nn.Module): |
| 94 | """A wrapper around a conv layer that behaves like a BaseBackbone.""" |
| 95 | |
| 96 | def __init__(self, dim_in: int, dim_out: int, kernel_size: int, stride_out: int): |
| 97 | """Initialize SkipConvBackbone.""" |
| 98 | super().__init__() |
| 99 | self.stride_out = stride_out |
| 100 | if stride_out == 1 and kernel_size != 1: |
| 101 | raise ValueError("We only support kernel_size = 1 if stride_out is 1.") |
| 102 | padding: int = (kernel_size - 1) // 2 |
| 103 | self.conv = nn.Conv2d( |
| 104 | dim_in, dim_out, kernel_size=kernel_size, stride=stride_out, padding=padding |
| 105 | ) |
| 106 | |
| 107 | def forward( |
| 108 | self, |
| 109 | input_features: torch.Tensor, |
| 110 | encodings: list[torch.Tensor] | None = None, |
| 111 | ) -> ImageFeatures: |
| 112 | """Apply SkipConvBackbone to image.""" |
| 113 | output = self.conv(input_features) |
| 114 | return ImageFeatures( |
| 115 | texture_features=output, |
| 116 | geometry_features=output, |
| 117 | ) |
| 118 | |
| 119 | @property |
| 120 | def stride(self) -> int: |
| 121 | """Effective downsampling stride.""" |
| 122 | return self.stride_out |
| 123 | |
| 124 | |
| 125 | class GaussianDensePredictionTransformer(nn.Module): |
no outgoing calls
no test coverage detected