288 lines
11 KiB
Python
288 lines
11 KiB
Python
# Copyright (c) OpenMMLab. All rights reserved.
|
|
import os
|
|
import shutil
|
|
import sys
|
|
import warnings
|
|
from unittest.mock import MagicMock
|
|
|
|
import numpy as np
|
|
import pytest
|
|
import torch
|
|
|
|
from mmengine import Config
|
|
from mmengine.fileio import load
|
|
from mmengine.registry import VISBACKENDS
|
|
from mmengine.visualization import (LocalVisBackend, MLflowVisBackend,
|
|
TensorboardVisBackend, WandbVisBackend)
|
|
|
|
|
|
class TestLocalVisBackend:
|
|
|
|
def test_init(self):
|
|
# 'config_save_file' format must be py
|
|
with pytest.raises(AssertionError):
|
|
LocalVisBackend('temp_dir', config_save_file='a.txt')
|
|
|
|
# 'scalar_save_file' format must be json
|
|
with pytest.raises(AssertionError):
|
|
LocalVisBackend('temp_dir', scalar_save_file='a.yaml')
|
|
|
|
local_vis_backend = VISBACKENDS.build(
|
|
dict(type='LocalVisBackend', save_dir='temp_dir'))
|
|
assert isinstance(local_vis_backend, LocalVisBackend)
|
|
|
|
def test_experiment(self):
|
|
local_vis_backend = LocalVisBackend('temp_dir')
|
|
assert local_vis_backend.experiment == local_vis_backend
|
|
|
|
def test_add_config(self):
|
|
cfg = Config(dict(a=1, b=dict(b1=[0, 1])))
|
|
local_vis_backend = LocalVisBackend('temp_dir')
|
|
local_vis_backend.add_config(cfg)
|
|
assert os.path.exists(local_vis_backend._config_save_file)
|
|
shutil.rmtree('temp_dir')
|
|
|
|
def test_add_image(self):
|
|
image = np.random.randint(0, 256, size=(10, 10, 3))
|
|
local_vis_backend = LocalVisBackend('temp_dir')
|
|
|
|
# image must be in np.uint8 format
|
|
with pytest.raises(AssertionError):
|
|
local_vis_backend.add_image('img', image)
|
|
|
|
local_vis_backend.add_image('img', image.astype(np.uint8))
|
|
assert os.path.exists(
|
|
os.path.join(local_vis_backend._img_save_dir, 'img_0.png'))
|
|
local_vis_backend.add_image('img', image.astype(np.uint8), step=2)
|
|
assert os.path.exists(
|
|
os.path.join(local_vis_backend._img_save_dir, 'img_2.png'))
|
|
shutil.rmtree('temp_dir')
|
|
|
|
def test_add_scalar(self):
|
|
local_vis_backend = LocalVisBackend('temp_dir')
|
|
local_vis_backend.add_scalar('map', 0.9)
|
|
out_dict = load(local_vis_backend._scalar_save_file, 'json')
|
|
assert out_dict == {'map': 0.9, 'step': 0}
|
|
shutil.rmtree('temp_dir')
|
|
|
|
# test append mode
|
|
local_vis_backend = LocalVisBackend('temp_dir')
|
|
local_vis_backend.add_scalar('map', 1, step=0)
|
|
local_vis_backend.add_scalar('map', 0.95, step=1)
|
|
# local_vis_backend.add_scalar('map', torch.IntTensor(1), step=2)
|
|
local_vis_backend.add_scalar('map', np.array(0.9), step=2)
|
|
with open(local_vis_backend._scalar_save_file) as f:
|
|
out_dict = f.read()
|
|
assert out_dict == '{"map": 1, "step": 0}\n' \
|
|
'{"map": 0.95, "step": 1}\n' \
|
|
'{"map": 0.9, "step": 2}\n'
|
|
shutil.rmtree('temp_dir')
|
|
|
|
local_vis_backend = LocalVisBackend('temp_dir')
|
|
local_vis_backend.add_scalar('map', torch.tensor(1.))
|
|
assert os.path.exists(local_vis_backend._scalar_save_file)
|
|
shutil.rmtree('temp_dir')
|
|
|
|
def test_add_scalars(self):
|
|
input_dict = {'map': 0.7, 'acc': 0.9}
|
|
local_vis_backend = LocalVisBackend('temp_dir')
|
|
local_vis_backend.add_scalars(input_dict)
|
|
assert input_dict == {'map': 0.7, 'acc': 0.9}
|
|
out_dict = load(local_vis_backend._scalar_save_file, 'json')
|
|
assert out_dict == {'map': 0.7, 'acc': 0.9, 'step': 0}
|
|
|
|
# test append mode
|
|
local_vis_backend.add_scalars({'map': 0.8, 'acc': 0.8}, step=1)
|
|
with open(local_vis_backend._scalar_save_file) as f:
|
|
out_dict = f.read()
|
|
assert out_dict == '{"map": 0.7, "acc": 0.9, ' \
|
|
'"step": 0}\n{"map": 0.8, "acc": 0.8, "step": 1}\n'
|
|
|
|
# test file_path
|
|
local_vis_backend.add_scalars(input_dict, file_path='temp.json')
|
|
assert os.path.exists(local_vis_backend._scalar_save_file)
|
|
assert os.path.exists(
|
|
os.path.join(local_vis_backend._save_dir, 'temp.json'))
|
|
|
|
# file_path and scalar_save_file cannot be the same
|
|
with pytest.raises(AssertionError):
|
|
local_vis_backend.add_scalars(input_dict, file_path='scalars.json')
|
|
|
|
shutil.rmtree('temp_dir')
|
|
|
|
|
|
class TestTensorboardVisBackend:
|
|
sys.modules['torch.utils.tensorboard'] = MagicMock()
|
|
sys.modules['tensorboardX'] = MagicMock()
|
|
|
|
def test_init(self):
|
|
TensorboardVisBackend('temp_dir')
|
|
VISBACKENDS.build(
|
|
dict(type='TensorboardVisBackend', save_dir='temp_dir'))
|
|
|
|
def test_experiment(self):
|
|
tensorboard_vis_backend = TensorboardVisBackend('temp_dir')
|
|
assert (tensorboard_vis_backend.experiment ==
|
|
tensorboard_vis_backend._tensorboard)
|
|
shutil.rmtree('temp_dir')
|
|
|
|
def test_add_config(self):
|
|
cfg = Config(dict(a=1, b=dict(b1=[0, 1])))
|
|
tensorboard_vis_backend = TensorboardVisBackend('temp_dir')
|
|
tensorboard_vis_backend.add_config(cfg)
|
|
shutil.rmtree('temp_dir')
|
|
|
|
def test_add_image(self):
|
|
image = np.random.randint(0, 256, size=(10, 10, 3)).astype(np.uint8)
|
|
tensorboard_vis_backend = TensorboardVisBackend('temp_dir')
|
|
tensorboard_vis_backend.add_image('img', image)
|
|
tensorboard_vis_backend.add_image('img', image, step=2)
|
|
shutil.rmtree('temp_dir')
|
|
|
|
def test_add_scalar(self):
|
|
tensorboard_vis_backend = TensorboardVisBackend('temp_dir')
|
|
tensorboard_vis_backend.add_scalar('map', 0.9)
|
|
# test append mode
|
|
tensorboard_vis_backend.add_scalar('map', 0.9, step=0)
|
|
tensorboard_vis_backend.add_scalar('map', 0.95, step=1)
|
|
# test with numpy
|
|
with warnings.catch_warnings(record=True) as record:
|
|
tensorboard_vis_backend.add_scalar('map', np.array(0.9), step=0)
|
|
tensorboard_vis_backend.add_scalar('map', np.array(0.95), step=1)
|
|
tensorboard_vis_backend.add_scalar('map', np.array(9), step=0)
|
|
tensorboard_vis_backend.add_scalar('map', np.array(95), step=1)
|
|
tensorboard_vis_backend.add_scalar('map', np.array([9])[0], step=0)
|
|
tensorboard_vis_backend.add_scalar(
|
|
'map', np.array([95])[0], step=1)
|
|
assert len(record) == 0
|
|
# test with tensor
|
|
tensorboard_vis_backend.add_scalar('map', torch.tensor(0.9), step=0)
|
|
tensorboard_vis_backend.add_scalar('map', torch.tensor(0.95), step=1)
|
|
# Unprocessable data will output a warning message
|
|
with pytest.warns(Warning):
|
|
tensorboard_vis_backend.add_scalar('map', [0.95])
|
|
shutil.rmtree('temp_dir')
|
|
|
|
def test_add_scalars(self):
|
|
tensorboard_vis_backend = TensorboardVisBackend('temp_dir')
|
|
# The step value must be passed through the parameter
|
|
with pytest.raises(AssertionError):
|
|
tensorboard_vis_backend.add_scalars({
|
|
'map': 0.7,
|
|
'acc': 0.9,
|
|
'step': 1
|
|
})
|
|
|
|
# Unprocessable data will output a warning message
|
|
with pytest.warns(Warning):
|
|
tensorboard_vis_backend.add_scalars({
|
|
'map': [1, 2],
|
|
})
|
|
|
|
input_dict = {'map': 0.7, 'acc': 0.9}
|
|
tensorboard_vis_backend.add_scalars(input_dict)
|
|
# test append mode
|
|
tensorboard_vis_backend.add_scalars({'map': 0.8, 'acc': 0.8}, step=1)
|
|
shutil.rmtree('temp_dir')
|
|
|
|
def test_close(self):
|
|
tensorboard_vis_backend = TensorboardVisBackend('temp_dir')
|
|
tensorboard_vis_backend._init_env()
|
|
tensorboard_vis_backend.close()
|
|
shutil.rmtree('temp_dir')
|
|
|
|
|
|
class TestWandbVisBackend:
|
|
sys.modules['wandb'] = MagicMock()
|
|
sys.modules['wandb.run'] = MagicMock()
|
|
|
|
def test_init(self):
|
|
WandbVisBackend('temp_dir')
|
|
VISBACKENDS.build(dict(type='WandbVisBackend', save_dir='temp_dir'))
|
|
|
|
def test_experiment(self):
|
|
wandb_vis_backend = WandbVisBackend('temp_dir')
|
|
assert wandb_vis_backend.experiment == wandb_vis_backend._wandb
|
|
shutil.rmtree('temp_dir')
|
|
|
|
def test_add_config(self):
|
|
cfg = Config(dict(a=1, b=dict(b1=[0, 1])))
|
|
wandb_vis_backend = WandbVisBackend('temp_dir', log_code_name='code')
|
|
_wandb = wandb_vis_backend.experiment
|
|
_wandb.run.dir = 'temp_dir'
|
|
wandb_vis_backend.add_config(cfg)
|
|
shutil.rmtree('temp_dir')
|
|
|
|
def test_add_image(self):
|
|
image = np.random.randint(0, 256, size=(10, 10, 3)).astype(np.uint8)
|
|
wandb_vis_backend = WandbVisBackend('temp_dir')
|
|
wandb_vis_backend.add_image('img', image)
|
|
wandb_vis_backend.add_image('img', image)
|
|
shutil.rmtree('temp_dir')
|
|
|
|
def test_add_scalar(self):
|
|
wandb_vis_backend = WandbVisBackend('temp_dir')
|
|
wandb_vis_backend.add_scalar('map', 0.9)
|
|
# test append mode
|
|
wandb_vis_backend.add_scalar('map', 0.9)
|
|
wandb_vis_backend.add_scalar('map', 0.95)
|
|
shutil.rmtree('temp_dir')
|
|
|
|
def test_add_scalars(self):
|
|
wandb_vis_backend = WandbVisBackend('temp_dir')
|
|
input_dict = {'map': 0.7, 'acc': 0.9}
|
|
wandb_vis_backend.add_scalars(input_dict)
|
|
# test append mode
|
|
wandb_vis_backend.add_scalars({'map': 0.8, 'acc': 0.8})
|
|
wandb_vis_backend.add_scalars({'map': [0.8], 'acc': 0.8})
|
|
shutil.rmtree('temp_dir')
|
|
|
|
def test_close(self):
|
|
wandb_vis_backend = WandbVisBackend('temp_dir')
|
|
wandb_vis_backend._init_env()
|
|
wandb_vis_backend.close()
|
|
shutil.rmtree('temp_dir')
|
|
|
|
|
|
class TestMLflowVisBackend:
|
|
|
|
def test_init(self):
|
|
MLflowVisBackend('temp_dir')
|
|
VISBACKENDS.build(dict(type='MLflowVisBackend', save_dir='temp_dir'))
|
|
|
|
def test_experiment(self):
|
|
mlflow_vis_backend = MLflowVisBackend('temp_dir')
|
|
assert mlflow_vis_backend.experiment == mlflow_vis_backend._mlflow
|
|
|
|
def test_add_config(self):
|
|
cfg = Config(dict(a=1, b=dict(b1=[0, 1])))
|
|
mlflow_vis_backend = MLflowVisBackend('temp_dir')
|
|
mlflow_vis_backend.add_config(cfg)
|
|
|
|
def test_add_image(self):
|
|
image = np.random.randint(0, 256, size=(10, 10, 3)).astype(np.uint8)
|
|
mlflow_vis_backend = MLflowVisBackend('temp_dir')
|
|
mlflow_vis_backend.add_image('img.png', image)
|
|
|
|
def test_add_scalar(self):
|
|
mlflow_vis_backend = MLflowVisBackend('temp_dir')
|
|
mlflow_vis_backend.add_scalar('map', 0.9)
|
|
# test append mode
|
|
mlflow_vis_backend.add_scalar('map', 0.9)
|
|
mlflow_vis_backend.add_scalar('map', 0.95)
|
|
|
|
def test_add_scalars(self):
|
|
mlflow_vis_backend = MLflowVisBackend('temp_dir')
|
|
input_dict = {'map': 0.7, 'acc': 0.9}
|
|
mlflow_vis_backend.add_scalars(input_dict)
|
|
# test append mode
|
|
mlflow_vis_backend.add_scalars({'map': 0.8, 'acc': 0.8})
|
|
|
|
def test_close(self):
|
|
cfg = Config(dict(work_dir='temp_dir'))
|
|
mlflow_vis_backend = MLflowVisBackend('temp_dir')
|
|
mlflow_vis_backend._init_env()
|
|
mlflow_vis_backend.add_config(cfg)
|
|
mlflow_vis_backend.close()
|
|
shutil.rmtree('temp_dir')
|