| 159 | self.opt.load_state_dict(state_dict) |
| 160 | |
| 161 | def run_loop(self): |
| 162 | i = 0 |
| 163 | |
| 164 | while ( |
| 165 | not self.lr_anneal_steps |
| 166 | or self.step + self.resume_step < self.lr_anneal_steps |
| 167 | ): |
| 168 | if self.dataset=='brats': |
| 169 | try: |
| 170 | batch, cond, label = next(self.iterdatal) |
| 171 | except: |
| 172 | self.iterdatal = iter(self.datal) |
| 173 | batch, cond, label, _, _ = next(self.iterdatal) |
| 174 | elif self.dataset=='chexpert': |
| 175 | batch, cond = next(self.datal) |
| 176 | cond.pop("path", None) |
| 177 | |
| 178 | self.run_step(batch, cond) |
| 179 | |
| 180 | if self.step % self.log_interval == 0: |
| 181 | logger.dumpkvs() |
| 182 | if self.step % self.save_interval == 0: |
| 183 | self.save() |
| 184 | # Run for a finite amount of time in integration tests. |
| 185 | if os.environ.get("DIFFUSION_TRAINING_TEST", "") and self.step > 0: |
| 186 | return |
| 187 | self.step += 1 |
| 188 | # Save the last checkpoint if it wasn't already saved. |
| 189 | if (self.step - 1) % self.save_interval != 0: |
| 190 | self.save() |
| 191 | |
| 192 | def run_step(self, batch, cond): |
| 193 | lossmse, sample = self.forward_backward(batch, cond) |