iti/HandMesh

1import os2import time3import torch4import json5from glob import glob6 7 8class Writer:9    def __init__(self, args=None):10        self.args = args11 12        if self.args is not None:13            tmp_log_list = glob(os.path.join(args.out_dir, 'log*'))14            if len(tmp_log_list) == 0:15                self.log_file = os.path.join(16                    args.out_dir, 'log_{:s}.txt'.format(17                        time.strftime('%Y-%m-%d-%H-%M-%S', time.localtime())))18            else:19                self.log_file = tmp_log_list[0]20 21    def print_str(self, info):22        with open(self.log_file, 'a') as log_file:23            log_file.write('{:s}\n'.format(str(info)))24        print(info)25 26    def print_info(self, info):27        message = 'Epoch: {}/{}, Duration: {:.3f}s, Train Loss: {:.4f}' \28                .format(info['current_epoch'], info['epochs'], info['t_duration'], info['train_loss'])29        with open(self.log_file, 'a') as log_file:30            log_file.write('{:s}\n'.format(message))31        # print(message)32 33    def print_step(self, info):34        message = 'Epoch: {}, Total_step: {}, Duration: {:.3f}s, Train Loss: {:.4f}, Lr: {:.6f}' \35            .format(info['epoch'], info['total_step'], info['step_duration'], info['train_loss'], info['lr'])36        with open(self.log_file, 'a') as log_file:37            log_file.write('{:s}\n'.format(message))38        # print(message)39 40    def save_checkpoint(self, model, optimizer, scheduler, epoch, best=False, last=False):41        if best:42            save_path = os.path.join(self.args.checkpoints_dir, 'checkpoint_best.pt')43        elif last:44            save_path = os.path.join(self.args.checkpoints_dir, 'checkpoint_last.pt')45        else:46            save_path = os.path.join(self.args.checkpoints_dir, 'checkpoint_{:03d}.pt'.format(epoch))47        torch.save(48            {49                'epoch': epoch,50                'model_state_dict': model.state_dict(),51                'optimizer_state_dict': optimizer.state_dict(),52                'scheduler_state_dict': scheduler.state_dict(),53            }, save_path)54