| 168 | self.optimizers['optimizer_Dis'].step() |
| 169 | |
| 170 | def test_iter(self, metrics=None): |
| 171 | if not self.is_train: |
| 172 | self.is_train = True |
| 173 | self.setup_net_parallel() |
| 174 | |
| 175 | self.nets['kp_detector'].eval() |
| 176 | self.nets['generator'].eval() |
| 177 | with paddle.no_grad(): |
| 178 | kp_source = self.nets['kp_detector'](self.input_data['video'][:, :, |
| 179 | 0]) |
| 180 | for frame_idx in range(self.input_data['video'].shape[2]): |
| 181 | source = self.input_data['video'][:, :, 0] |
| 182 | driving = self.input_data['video'][:, :, frame_idx] |
| 183 | kp_driving = self.nets['kp_detector'](driving) |
| 184 | out = self.nets['generator'](source, |
| 185 | kp_source=kp_source, |
| 186 | kp_driving=kp_driving) |
| 187 | out.update({'kp_source': kp_source, 'kp_driving': kp_driving}) |
| 188 | loss = paddle.abs(out['prediction'] - |
| 189 | driving).mean().cpu().numpy() |
| 190 | self.test_loss.append(loss) |
| 191 | self.visual_items['driving_source_gen'] = self.visualizer.visualize( |
| 192 | driving, source, out) |
| 193 | print("Reconstruction loss: %s" % np.mean(self.test_loss)) |
| 194 | self.nets['kp_detector'].train() |
| 195 | self.nets['generator'].train() |
| 196 | |
| 197 | class InferGenerator(paddle.nn.Layer): |
| 198 | def set_generator(self, generator): |