Initialize UNet Encoder. Args: dim_in: The number of input channels. width: Width multiplicator of intermediate layers or the width list of all layers. steps: The number of downsampling steps. norm_type: Which kind of normalization layer to us
(
self,
dim_in: int,
width: List[int] | int,
steps: int = 6,
norm_type: NormLayerName = "group_norm",
norm_num_groups=8,
blocks_per_layer=2,
)
| 24 | """Encoder of UNet model.""" |
| 25 | |
| 26 | def __init__( |
| 27 | self, |
| 28 | dim_in: int, |
| 29 | width: List[int] | int, |
| 30 | steps: int = 6, |
| 31 | norm_type: NormLayerName = "group_norm", |
| 32 | norm_num_groups=8, |
| 33 | blocks_per_layer=2, |
| 34 | ) -> None: |
| 35 | """Initialize UNet Encoder. |
| 36 | |
| 37 | Args: |
| 38 | dim_in: The number of input channels. |
| 39 | width: Width multiplicator of intermediate layers or the width list of all layers. |
| 40 | steps: The number of downsampling steps. |
| 41 | norm_type: Which kind of normalization layer to use. |
| 42 | norm_num_groups: How many groups to use for group norm (if relevant). |
| 43 | blocks_per_layer: How many residual blocks per layer to use. |
| 44 | """ |
| 45 | super().__init__() |
| 46 | |
| 47 | if blocks_per_layer < 1: |
| 48 | raise ValueError("blocks_per_layer must be greater or equal to one.") |
| 49 | |
| 50 | self.dim_in = dim_in |
| 51 | self.width = width |
| 52 | self.num_steps = steps |
| 53 | |
| 54 | self.convs_down = nn.ModuleList() |
| 55 | |
| 56 | self.output_dims: list[int] |
| 57 | # If only one number is specified, we assume each layer will double the channel dimension. |
| 58 | if isinstance(width, int): |
| 59 | self.output_dims = [width << i for i in range(0, steps + 1)] |
| 60 | else: |
| 61 | if len(width) != (steps + 1): |
| 62 | raise ValueError("Length of width should match the steps for UNetEncoder.") |
| 63 | self.output_dims = width |
| 64 | |
| 65 | self.conv_in = nn.Sequential( |
| 66 | nn.Conv2d(self.dim_in, self.output_dims[0], 3, stride=1, padding=1), |
| 67 | norm_layer_2d(self.output_dims[0], norm_type, num_groups=norm_num_groups), |
| 68 | nn.ReLU(), |
| 69 | ) |
| 70 | |
| 71 | for i_step in range(steps): |
| 72 | input_width = self.output_dims[i_step] |
| 73 | current_width = self.output_dims[i_step + 1] |
| 74 | convs_down_i = nn.Sequential( |
| 75 | nn.AvgPool2d(2, stride=2), |
| 76 | residual_block_2d( |
| 77 | input_width, |
| 78 | current_width, |
| 79 | norm_type=norm_type, |
| 80 | norm_num_groups=norm_num_groups, |
| 81 | ), |
| 82 | *[ |
| 83 | residual_block_2d( |
nothing calls this directly
no test coverage detected