| 53 | return NetworksFactory.get_by_name('discriminator_wasserstein_gan', c_dim=self._opt.cond_nc) |
| 54 | |
| 55 | def _init_train_vars(self): |
| 56 | self._current_lr_G = self._opt.lr_G |
| 57 | self._current_lr_D = self._opt.lr_D |
| 58 | |
| 59 | # initialize optimizers |
| 60 | self._optimizer_G = torch.optim.Adam(self._G.parameters(), lr=self._current_lr_G, |
| 61 | betas=[self._opt.G_adam_b1, self._opt.G_adam_b2]) |
| 62 | self._optimizer_D = torch.optim.Adam(self._D.parameters(), lr=self._current_lr_D, |
| 63 | betas=[self._opt.D_adam_b1, self._opt.D_adam_b2]) |
| 64 | |
| 65 | def _init_prefetch_inputs(self): |
| 66 | self._input_real_img = self._Tensor(self._opt.batch_size, 3, self._opt.image_size, self._opt.image_size) |