File size: 485 Bytes
affcd23 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 |
from enum import Enum
from .tensorboard import TensorboardWriter
from .wandb import WanDBWriter
class VisualizerBackendType(str, Enum):
tensorboard = "tensorboard"
wandb = "wandb"
def get_visualizer(config, logger, backend: VisualizerBackendType):
if backend == VisualizerBackendType.tensorboard:
return TensorboardWriter(config.log_dir, logger, True)
if backend == VisualizerBackendType.wandb:
return WanDBWriter(config, logger)
return None
|