(self)
| 169 | |
| 170 | |
| 171 | def train(self): |
| 172 | optimizer = optim.AdamW(self.network.parameters(), lr=self.LR, weight_decay=1e-4) |
| 173 | iteration = self.continue_training_steps |
| 174 | save_img_path = init__result_Dir() |
| 175 | print('Starting Training', f"Step is {self.num_timesteps}") |
| 176 | |
| 177 | while iteration < self.iteration_max: |
| 178 | |
| 179 | tq = tqdm(self.dataloader_train) |
| 180 | |
| 181 | for img, gt, name in tq: |
| 182 | tq.set_description(f'Iteration {iteration} / {self.iteration_max}') |
| 183 | self.network.train() |
| 184 | optimizer.zero_grad() |
| 185 | |
| 186 | t = torch.randint(0, self.num_timesteps, (img.shape[0],)).long().to(self.device) |
| 187 | init_predict, noise_pred, noisy_image, noise_ref = self.network(gt.to(self.device), img.to(self.device), |
| 188 | t, self.diffusion) |
| 189 | if self.pre_ori == 'True': |
| 190 | if self.high_low_freq == 'True': |
| 191 | residual_high = self.high_filter(gt.to(self.device) - init_predict) |
| 192 | ddpm_loss = 2*self.loss(self.high_filter(noise_pred), residual_high) + self.loss(noise_pred, gt.to(self.device) - init_predict) |
| 193 | else: |
| 194 | ddpm_loss = self.loss(noise_pred, gt.to(self.device) - init_predict) |
| 195 | else: |
| 196 | ddpm_loss = self.loss(noise_pred, noise_ref.to(self.device)) |
| 197 | if self.high_low_freq == 'True': |
| 198 | low_high_loss = self.loss(init_predict, gt.to(self.device)) |
| 199 | low_freq_loss = self.loss(init_predict - self.high_filter(init_predict), gt.to(self.device) - self.high_filter(gt.to(self.device))) |
| 200 | pixel_loss = low_high_loss + 2*low_freq_loss |
| 201 | else: |
| 202 | pixel_loss = self.loss(init_predict, gt.to(self.device)) |
| 203 | |
| 204 | loss = ddpm_loss + self.beta_loss * (pixel_loss) / self.num_timesteps |
| 205 | loss.backward() |
| 206 | optimizer.step() |
| 207 | if self.high_low_freq == 'True': |
| 208 | tq.set_postfix(loss=loss.item(), high_freq_ddpm_loss=ddpm_loss.item(), low_freq_pixel_loss=low_freq_loss.item(), pixel_loss=low_high_loss.item()) |
| 209 | else: |
| 210 | tq.set_postfix(loss=loss.item(), ddpm_loss=ddpm_loss.item(), pixel_loss=pixel_loss.item()) |
| 211 | if iteration % 500 == 0: |
| 212 | if not os.path.exists(save_img_path): |
| 213 | os.makedirs(save_img_path) |
| 214 | img_save = torch.cat([img, gt, init_predict.cpu()], dim=3) |
| 215 | if self.pre_ori == 'True': |
| 216 | if self.high_low_freq == 'True': |
| 217 | img_save = torch.cat([img, gt, init_predict.cpu(), noise_pred.cpu() + self.high_filter(init_predict).cpu(), noise_pred.cpu() + init_predict.cpu()], dim=3) |
| 218 | else: |
| 219 | img_save = torch.cat([img, gt, init_predict.cpu(), noise_pred.cpu() + init_predict.cpu()], dim=3) |
| 220 | save_image(img_save, os.path.join( |
| 221 | save_img_path, f"{iteration}.png"), nrow=4) |
| 222 | iteration += 1 |
| 223 | if self.EMA_or_not == 'True': |
| 224 | if iteration % self.ema_every == 0 and iteration > self.start_ema: |
| 225 | print('EMA update') |
| 226 | self.EMA.update_model_average(self.ema_model, self.network) |
| 227 | |
| 228 | if iteration % self.save_model_every == 0: |
no test coverage detected