Flax implementation of 2D Resnet Block. Args: in_channels (`int`): Input channels out_channels (`int`): Output channels dropout (:obj:`float`, *optional*, defaults to 0.0): Dropout rate groups (:obj:`int`, *optional*, defa
| 123 | |
| 124 | |
| 125 | class FlaxResnetBlock2D(nn.Module): |
| 126 | """ |
| 127 | Flax implementation of 2D Resnet Block. |
| 128 | |
| 129 | Args: |
| 130 | in_channels (`int`): |
| 131 | Input channels |
| 132 | out_channels (`int`): |
| 133 | Output channels |
| 134 | dropout (:obj:`float`, *optional*, defaults to 0.0): |
| 135 | Dropout rate |
| 136 | groups (:obj:`int`, *optional*, defaults to `32`): |
| 137 | The number of groups to use for group norm. |
| 138 | use_nin_shortcut (:obj:`bool`, *optional*, defaults to `None`): |
| 139 | Whether to use `nin_shortcut`. This activates a new layer inside ResNet block |
| 140 | dtype (:obj:`jnp.dtype`, *optional*, defaults to jnp.float32): |
| 141 | Parameters `dtype` |
| 142 | """ |
| 143 | |
| 144 | in_channels: int |
| 145 | out_channels: int = None |
| 146 | dropout: float = 0.0 |
| 147 | groups: int = 32 |
| 148 | use_nin_shortcut: bool = None |
| 149 | dtype: jnp.dtype = jnp.float32 |
| 150 | |
| 151 | def setup(self): |
| 152 | out_channels = self.in_channels if self.out_channels is None else self.out_channels |
| 153 | |
| 154 | self.norm1 = nn.GroupNorm(num_groups=self.groups, epsilon=1e-6) |
| 155 | self.conv1 = nn.Conv( |
| 156 | out_channels, |
| 157 | kernel_size=(3, 3), |
| 158 | strides=(1, 1), |
| 159 | padding=((1, 1), (1, 1)), |
| 160 | dtype=self.dtype, |
| 161 | ) |
| 162 | |
| 163 | self.norm2 = nn.GroupNorm(num_groups=self.groups, epsilon=1e-6) |
| 164 | self.dropout_layer = nn.Dropout(self.dropout) |
| 165 | self.conv2 = nn.Conv( |
| 166 | out_channels, |
| 167 | kernel_size=(3, 3), |
| 168 | strides=(1, 1), |
| 169 | padding=((1, 1), (1, 1)), |
| 170 | dtype=self.dtype, |
| 171 | ) |
| 172 | |
| 173 | use_nin_shortcut = self.in_channels != out_channels if self.use_nin_shortcut is None else self.use_nin_shortcut |
| 174 | |
| 175 | self.conv_shortcut = None |
| 176 | if use_nin_shortcut: |
| 177 | self.conv_shortcut = nn.Conv( |
| 178 | out_channels, |
| 179 | kernel_size=(1, 1), |
| 180 | strides=(1, 1), |
| 181 | padding="VALID", |
| 182 | dtype=self.dtype, |