Spaces:
Running on Zero
Running on Zero
File size: 2,191 Bytes
2680bd5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 | import os
import wandb
from omegaconf import OmegaConf
from .. import dist as dist_utils
class TrainPlatform:
def __init__(self, save_dir):
pass
def report_scalar(self, name, value, iteration, group_name=None):
pass
def report_args(self, args, name):
pass
def close(self):
pass
class TensorboardPlatform(TrainPlatform):
def __init__(self, save_dir):
from torch.utils.tensorboard import SummaryWriter
self.writer = SummaryWriter(log_dir=save_dir)
def report_scalar(self, name, value, iteration, group_name=None):
self.writer.add_scalar(f"{group_name}/{name}" if group_name else name, value, iteration)
def close(self):
self.writer.close()
class WandbPlatform(TrainPlatform):
def __init__(self, save_dir):
if not dist_utils.is_main_process():
return
self.run = wandb.init(
dir=save_dir,
project='BihandMotionGeneration_FINAL'
)
self.run.log_code(name='source_code')
def report_scalar(self, name, value, iteration, group_name=None):
if not dist_utils.is_main_process():
return
log_key = f"{group_name}/{name}" if group_name else name
wandb.log({log_key:value}, step=iteration)
def report_args(self, args, name):
if not dist_utils.is_main_process():
return
try:
args_dict = OmegaConf.to_container(args, resolve=True)
wandb.config.update({name: args_dict})
except Exception as e:
raise ValueError(f'WandbPlatform: Could not log args {name}. Error: {e}')
def report_video(self, name, video_path, video_format, group_name=None):
if not dist_utils.is_main_process():
return
log_key = f"{group_name}/{name}" if group_name else name
wandb.log({
log_key: wandb.Video(video_path, format=video_format)
})
def close(self):
if not dist_utils.is_main_process():
return
if self.run:
wandb.finish()
self.run = None
class NoPlatform(TrainPlatform):
def __init__(self, save_dir):
pass
|