(self)
| 262 | self.vae_conv_final = self._make_final_conv(in_channels) |
| 263 | |
| 264 | def _prepare_vae_modules(self): |
| 265 | zoom = 2 ** (len(self.blocks_down) - 1) |
| 266 | v_filters = self.init_filters * zoom |
| 267 | total_elements = int(self.smallest_filters * np.prod(self.fc_insize)) |
| 268 | |
| 269 | self.vae_down = nn.Sequential( |
| 270 | get_norm_layer(name=self.norm, spatial_dims=self.spatial_dims, channels=v_filters), |
| 271 | self.act_mod, |
| 272 | get_conv_layer(self.spatial_dims, v_filters, self.smallest_filters, stride=2, bias=True), |
| 273 | get_norm_layer(name=self.norm, spatial_dims=self.spatial_dims, channels=self.smallest_filters), |
| 274 | self.act_mod, |
| 275 | ) |
| 276 | self.vae_fc1 = nn.Linear(total_elements, self.vae_nz) |
| 277 | self.vae_fc2 = nn.Linear(total_elements, self.vae_nz) |
| 278 | self.vae_fc3 = nn.Linear(self.vae_nz, total_elements) |
| 279 | |
| 280 | self.vae_fc_up_sample = nn.Sequential( |
| 281 | get_conv_layer(self.spatial_dims, self.smallest_filters, v_filters, kernel_size=1), |
| 282 | get_upsample_layer(self.spatial_dims, v_filters, upsample_mode=self.upsample_mode), |
| 283 | get_norm_layer(name=self.norm, spatial_dims=self.spatial_dims, channels=v_filters), |
| 284 | self.act_mod, |
| 285 | ) |
| 286 | |
| 287 | def _get_vae_loss(self, net_input: torch.Tensor, vae_input: torch.Tensor): |
| 288 | """ |
no test coverage detected