(self, current_iter)
| 81 | self.mask = data['mask'].to(self.device) if 'mask' in data else None |
| 82 | |
| 83 | def optimize_parameters(self, current_iter): |
| 84 | self.optimizer_g.zero_grad() |
| 85 | self.output = self.net_g(self.lq, self.mask) |
| 86 | |
| 87 | l_total = 0 |
| 88 | loss_dict = OrderedDict() |
| 89 | |
| 90 | if self.cri_pix: |
| 91 | l_pix = self.cri_pix(self.output, self.gt, self.mask) |
| 92 | l_total += l_pix |
| 93 | loss_dict['l_pix'] = l_pix |
| 94 | |
| 95 | if self.cri_perceptual: |
| 96 | l_percep = self.cri_perceptual(self.output, self.gt) |
| 97 | if l_percep is not None: |
| 98 | l_total += l_percep |
| 99 | loss_dict['l_percep'] = l_percep |
| 100 | |
| 101 | l_total.backward() |
| 102 | self.optimizer_g.step() |
| 103 | |
| 104 | self.log_dict = self.reduce_loss_dict(loss_dict) |
| 105 | |
| 106 | if self.ema_decay > 0: |
| 107 | self.model_ema(decay=self.ema_decay) |
| 108 | |
| 109 | def test(self): |
| 110 | if hasattr(self, 'net_g_ema'): |
nothing calls this directly
no test coverage detected