| 6 | |
| 7 | |
| 8 | class TBVisualizer: |
| 9 | def __init__(self, opt): |
| 10 | self._opt = opt |
| 11 | self._save_path = os.path.join(opt.checkpoints_dir, opt.name) |
| 12 | |
| 13 | self._log_path = os.path.join(self._save_path, 'loss_log2.txt') |
| 14 | self._tb_path = os.path.join(self._save_path, 'summary.json') |
| 15 | self._writer = SummaryWriter(self._save_path) |
| 16 | |
| 17 | with open(self._log_path, "a") as log_file: |
| 18 | now = time.strftime("%c") |
| 19 | log_file.write('================ Training Loss (%s) ================\n' % now) |
| 20 | |
| 21 | def __del__(self): |
| 22 | self._writer.close() |
| 23 | |
| 24 | def display_current_results(self, visuals, it, is_train, save_visuals=False): |
| 25 | for label, image_numpy in visuals.items(): |
| 26 | sum_name = '{}/{}'.format('Train' if is_train else 'Test', label) |
| 27 | self._writer.add_image(sum_name, image_numpy, it) |
| 28 | |
| 29 | if save_visuals: |
| 30 | util.save_image(image_numpy, |
| 31 | os.path.join(self._opt.checkpoints_dir, self._opt.name, |
| 32 | 'event_imgs', sum_name, '%08d.png' % it)) |
| 33 | |
| 34 | self._writer.export_scalars_to_json(self._tb_path) |
| 35 | |
| 36 | def plot_scalars(self, scalars, it, is_train): |
| 37 | for label, scalar in scalars.items(): |
| 38 | sum_name = '{}/{}'.format('Train' if is_train else 'Test', label) |
| 39 | self._writer.add_scalar(sum_name, scalar, it) |
| 40 | |
| 41 | def print_current_train_errors(self, epoch, i, iters_per_epoch, errors, t, visuals_were_stored): |
| 42 | log_time = time.strftime("[%d/%m/%Y %H:%M:%S]") |
| 43 | visuals_info = "v" if visuals_were_stored else "" |
| 44 | message = '%s (T%s, epoch: %d, it: %d/%d, t/smpl: %.3fs) ' % (log_time, visuals_info, epoch, i, iters_per_epoch, t) |
| 45 | for k, v in errors.items(): |
| 46 | message += '%s:%.3f ' % (k, v) |
| 47 | |
| 48 | print(message) |
| 49 | with open(self._log_path, "a") as log_file: |
| 50 | log_file.write('%s\n' % message) |
| 51 | |
| 52 | def print_current_validate_errors(self, epoch, errors, t): |
| 53 | log_time = time.strftime("[%d/%m/%Y %H:%M:%S]") |
| 54 | message = '%s (V, epoch: %d, time_to_val: %ds) ' % (log_time, epoch, t) |
| 55 | for k, v in errors.items(): |
| 56 | message += '%s:%.3f ' % (k, v) |
| 57 | |
| 58 | print(message) |
| 59 | with open(self._log_path, "a") as log_file: |
| 60 | log_file.write('%s\n' % message) |
| 61 | |
| 62 | def save_images(self, visuals): |
| 63 | for label, image_numpy in visuals.items(): |
| 64 | image_name = '%s.png' % label |
| 65 | save_path = os.path.join(self._save_path, "samples", image_name) |