| 61 | from model.blocks import GatedConv, GatedDeconv |
| 62 | |
| 63 | def save_target(child, out_subdir): |
| 64 | if not os.path.exists(out_subdir): |
| 65 | os.makedirs(out_subdir) |
| 66 | if isinstance(child, GatedConv): |
| 67 | target = child.gated_values[0] |
| 68 | elif isinstance(child, GatedDeconv): |
| 69 | target = child.conv.gated_values[0] |
| 70 | else: |
| 71 | raise ValueError('should be gated conv or gated deconv') |
| 72 | target = target.transpose(0, 1) |
| 73 | for t in range(target.shape[0]): |
| 74 | for c in range(target.shape[1]): |
| 75 | out_file = os.path.join(out_subdir, f'time{t:03d}_channel{c:04d}.png') |
| 76 | self.toPILImage(target[t, c: c + 1]).save(out_file) |
| 77 | |
| 78 | for key, child in self.model.generator.coarse_net.upsample_module.named_children(): |
| 79 | out_subdir = os.path.join(out_dir, f'upsample_{key}') |