()
| 136 | |
| 137 | |
| 138 | def main(): |
| 139 | # initialize gen and disc, note: discriminator should be called critic, |
| 140 | # according to WGAN paper (since it no longer outputs between [0, 1]) |
| 141 | # but really who cares.. |
| 142 | gen = Generator( |
| 143 | config.Z_DIM, config.IN_CHANNELS, img_channels=config.CHANNELS_IMG |
| 144 | ).to(config.DEVICE) |
| 145 | critic = Discriminator( |
| 146 | config.Z_DIM, config.IN_CHANNELS, img_channels=config.CHANNELS_IMG |
| 147 | ).to(config.DEVICE) |
| 148 | |
| 149 | # initialize optimizers and scalers for FP16 training |
| 150 | opt_gen = optim.Adam(gen.parameters(), lr=config.LEARNING_RATE, betas=(0.0, 0.99)) |
| 151 | opt_critic = optim.Adam( |
| 152 | critic.parameters(), lr=config.LEARNING_RATE, betas=(0.0, 0.99) |
| 153 | ) |
| 154 | scaler_critic = torch.cuda.amp.GradScaler() |
| 155 | scaler_gen = torch.cuda.amp.GradScaler() |
| 156 | |
| 157 | # for tensorboard plotting |
| 158 | writer = SummaryWriter(f"logs/gan1") |
| 159 | |
| 160 | if config.LOAD_MODEL: |
| 161 | load_checkpoint( |
| 162 | config.PRE_CHECKPOINT_GEN, gen, opt_gen, config.LEARNING_RATE, |
| 163 | ) |
| 164 | load_checkpoint( |
| 165 | config.PRE_CHECKPOINT_CRITIC, critic, opt_critic, config.LEARNING_RATE, |
| 166 | ) |
| 167 | |
| 168 | gen.train() |
| 169 | critic.train() |
| 170 | |
| 171 | tensorboard_step = 0 |
| 172 | # start at step that corresponds to img size that we set in config |
| 173 | #step = 5 |
| 174 | step = int(log2(config.START_TRAIN_AT_IMG_SIZE / 4)) |
| 175 | for num_epochs in config.PROGRESSIVE_EPOCHS[step:7]: |
| 176 | alpha = 1e-5 # start with very low alpha |
| 177 | loader, dataset = get_loader(4 * 2 ** step) # 4->0, 8->1, 16->2, 32->3, 64 -> 4 |
| 178 | print(f"Current image size: {4 * 2 ** step}") |
| 179 | |
| 180 | for epoch in range(num_epochs): |
| 181 | print(f"Epoch [{epoch+1}/{num_epochs}]") |
| 182 | tensorboard_step, alpha = train_fn( |
| 183 | critic, |
| 184 | gen, |
| 185 | loader, |
| 186 | dataset, |
| 187 | step, |
| 188 | alpha, |
| 189 | opt_critic, |
| 190 | opt_gen, |
| 191 | tensorboard_step, |
| 192 | writer, |
| 193 | scaler_gen, |
| 194 | scaler_critic, |
| 195 | ) |
no test coverage detected