(self, batch, cond)
| 196 | self.ddp_model.train() |
| 197 | |
| 198 | def forward_backward(self, batch, cond): |
| 199 | self.mp_trainer.zero_grad() |
| 200 | for i in range(0, batch.shape[0], self.microbatch): |
| 201 | # Eliminates the microbatch feature |
| 202 | assert i == 0 |
| 203 | assert self.microbatch == self.batch_size |
| 204 | micro = batch.to(dist_util.dev()) |
| 205 | micro_cond = cond |
| 206 | # micro_cond = {k: v.to(dist_util.dev()) for k, v in cond.items()} |
| 207 | last_batch = (i + self.microbatch) >= batch.shape[0] |
| 208 | t, weights = self.schedule_sampler.sample(micro.shape[0], dist_util.dev()) |
| 209 | |
| 210 | compute_losses = functools.partial( |
| 211 | self.diffusion.training_losses, |
| 212 | self.ddp_model, |
| 213 | micro, |
| 214 | t, |
| 215 | model_kwargs=micro_cond, |
| 216 | ) |
| 217 | |
| 218 | if last_batch or not self.use_ddp: |
| 219 | losses = compute_losses() |
| 220 | else: |
| 221 | with self.ddp_model.no_sync(): |
| 222 | losses = compute_losses() |
| 223 | |
| 224 | if isinstance(self.schedule_sampler, LossAwareSampler): |
| 225 | self.schedule_sampler.update_with_local_losses( |
| 226 | t, losses["loss"].detach() |
| 227 | ) |
| 228 | |
| 229 | loss = (losses["loss"] * weights).mean() |
| 230 | self.mp_trainer.backward(loss) |
| 231 | |
| 232 | if self.step % 10 == 0: |
| 233 | self.log_loss_dict( |
| 234 | self.diffusion, t, {k: v * weights for k, v in losses.items()} |
| 235 | ) |
| 236 | |
| 237 | def _update_ema(self): |
| 238 | for rate, params in zip(self.ema_rate, self.ema_params): |
no test coverage detected