Initialize LearnedAlignment. Args: steps: Number of steps in the UNet. stride: Effective downsampling of the alignment module. base_width: Base width of the UNet. depth_decoder_features: Whether to use depth decoder features. depth
(
self,
steps: int = 4,
stride: int = 8,
base_width: int = 16,
depth_decoder_features: bool = False,
depth_decoder_dim: int = 256,
activation_type: math_utils.ActivationType = "exp",
)
| 44 | """Aligns tensors using a UNet.""" |
| 45 | |
| 46 | def __init__( |
| 47 | self, |
| 48 | steps: int = 4, |
| 49 | stride: int = 8, |
| 50 | base_width: int = 16, |
| 51 | depth_decoder_features: bool = False, |
| 52 | depth_decoder_dim: int = 256, |
| 53 | activation_type: math_utils.ActivationType = "exp", |
| 54 | ) -> None: |
| 55 | """Initialize LearnedAlignment. |
| 56 | |
| 57 | Args: |
| 58 | steps: Number of steps in the UNet. |
| 59 | stride: Effective downsampling of the alignment module. |
| 60 | base_width: Base width of the UNet. |
| 61 | depth_decoder_features: Whether to use depth decoder features. |
| 62 | depth_decoder_dim: Dimension of the depth decoder features. |
| 63 | activation_type: Activation type for the alignment output. |
| 64 | """ |
| 65 | super().__init__() |
| 66 | self.activation = math_utils.create_activation_pair(activation_type) |
| 67 | bias_value = self.activation.inverse(torch.tensor(1.0)) |
| 68 | |
| 69 | self.depth_decoder_features = depth_decoder_features |
| 70 | if depth_decoder_features: |
| 71 | dim_in = 2 + depth_decoder_dim |
| 72 | else: |
| 73 | dim_in = 2 |
| 74 | |
| 75 | def is_power_of_two(n: int) -> bool: |
| 76 | """Check if a number is a power of two.""" |
| 77 | if n <= 0: |
| 78 | return False |
| 79 | return (n & (n - 1)) == 0 |
| 80 | |
| 81 | if not is_power_of_two(stride): |
| 82 | raise ValueError(f"Stride {stride} is not a power of two.") |
| 83 | |
| 84 | steps_decoder = steps - int(math.log2(stride)) |
| 85 | if steps_decoder < 1: |
| 86 | raise ValueError(f"{steps_decoder} must be greater or equal to 1.") |
| 87 | widths = [min(base_width << i, 1024) for i in range(steps + 1)] |
| 88 | self.encoder = UNetEncoder(dim_in=dim_in, width=widths, steps=steps, norm_num_groups=4) |
| 89 | self.decoder = UNetDecoder( |
| 90 | dim_out=widths[0], width=widths, steps=steps_decoder, norm_num_groups=4 |
| 91 | ) |
| 92 | self.conv_out = nn.Conv2d(widths[0], 1, 1, bias=True) |
| 93 | nn.init.zeros_(self.conv_out.weight) |
| 94 | nn.init.constant_(self.conv_out.bias, bias_value) |
| 95 | |
| 96 | def forward( |
| 97 | self, |
nothing calls this directly
no test coverage detected