MCPcopy Index your code
hub / github.com/VisionXLab/OF-Diff / training_step

Method training_step

ldm/models/diffusion/ddpm.py:439–461  ·  view source on GitHub ↗
(self, batch, batch_idx)

Source from the content-addressed store, hash-verified

437 return loss, loss_dict
438
439 def training_step(self, batch, batch_idx):
440 for k in self.ucg_training:
441 p = self.ucg_training[k]["p"]
442 val = self.ucg_training[k]["val"]
443 if val is None:
444 val = ""
445 for i in range(len(batch[k])):
446 if self.ucg_prng.choice(2, p=[1 - p, p]):
447 batch[k][i] = val
448
449 loss, loss_dict = self.shared_step(batch)
450
451 self.log_dict(loss_dict, prog_bar=True,
452 logger=True, on_step=True, on_epoch=True)
453
454 self.log("global_step", self.global_step,
455 prog_bar=True, logger=True, on_step=True, on_epoch=False)
456
457 if self.use_scheduler:
458 lr = self.optimizers().param_groups[0]['lr']
459 self.log('lr_abs', lr, prog_bar=True, logger=True, on_step=True, on_epoch=False)
460
461 return loss
462
463 @torch.no_grad()
464 def validation_step(self, batch, batch_idx):

Callers

nothing calls this directly

Calls 1

shared_stepMethod · 0.95

Tested by

no test coverage detected