485 lines
17 KiB
Python
485 lines
17 KiB
Python
# Copyright (c) OpenMMLab. All rights reserved.
|
|
import os
|
|
import shutil
|
|
import sys
|
|
from unittest.mock import MagicMock, Mock, patch
|
|
|
|
import numpy as np
|
|
import pytest
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
from mmengine.fileio import load
|
|
from mmengine.registry import VISUALIZERS, WRITERS
|
|
from mmengine.visualization import (ComposedWriter, LocalWriter,
|
|
TensorboardWriter, WandbWriter)
|
|
|
|
|
|
def draw(self, image, gt_sample, pred_sample, show_gt=True, show_pred=True):
|
|
self.set_image(image)
|
|
|
|
|
|
class TestLocalWriter:
|
|
|
|
def test_init(self):
|
|
# visuailzer must be a dictionary or an instance
|
|
# of Visualizer and its subclasses
|
|
with pytest.raises(AssertionError):
|
|
LocalWriter('temp_dir', [dict(type='Visualizer')])
|
|
|
|
# 'params_save_file' format must be yaml
|
|
with pytest.raises(AssertionError):
|
|
LocalWriter('temp_dir', params_save_file='a.txt')
|
|
|
|
# 'scalar_save_file' format must be json
|
|
with pytest.raises(AssertionError):
|
|
LocalWriter('temp_dir', scalar_save_file='a.yaml')
|
|
|
|
local_writer = LocalWriter('temp_dir')
|
|
assert os.path.exists(local_writer._save_dir)
|
|
shutil.rmtree('temp_dir')
|
|
|
|
local_writer = WRITERS.build(
|
|
dict(
|
|
type='LocalWriter',
|
|
visualizer=dict(type='Visualizer'),
|
|
save_dir='temp_dir'))
|
|
assert os.path.exists(local_writer._save_dir)
|
|
shutil.rmtree('temp_dir')
|
|
|
|
def test_experiment(self):
|
|
local_writer = LocalWriter('temp_dir')
|
|
assert local_writer.experiment == local_writer
|
|
shutil.rmtree('temp_dir')
|
|
|
|
def test_add_params(self):
|
|
local_writer = LocalWriter('temp_dir')
|
|
|
|
# 'params_dict' must be dict
|
|
with pytest.raises(AssertionError):
|
|
local_writer.add_params(['lr', 0])
|
|
|
|
params_dict = dict(lr=0.1, wd=[1.0, 0.1, 0.001], mode='linear')
|
|
local_writer.add_params(params_dict)
|
|
out_dict = load(local_writer._params_save_file, 'yaml')
|
|
assert out_dict == params_dict
|
|
shutil.rmtree('temp_dir')
|
|
|
|
@patch('mmengine.visualization.visualizer.Visualizer.draw', draw)
|
|
def test_add_image(self):
|
|
image = np.random.randint(0, 256, size=(10, 10, 3)).astype(np.uint8)
|
|
|
|
# The visuailzer parameter must be set when
|
|
# the local_writer object is instantiated and
|
|
# the `add_image` method is called.
|
|
with pytest.raises(AssertionError):
|
|
local_writer = LocalWriter('temp_dir')
|
|
local_writer.add_image('img', image)
|
|
|
|
local_writer = LocalWriter('temp_dir', dict(type='Visualizer'))
|
|
local_writer.add_image('img', image)
|
|
assert os.path.exists(
|
|
os.path.join(local_writer._img_save_dir, 'img_0.png'))
|
|
|
|
bboxes = np.array([[1, 1, 2, 2], [1, 1.5, 1, 2.5]])
|
|
local_writer.visualizer.draw_bboxes(bboxes)
|
|
local_writer.add_image(
|
|
'img', local_writer.visualizer.get_image(), step=2)
|
|
assert os.path.exists(
|
|
os.path.join(local_writer._img_save_dir, 'img_2.png'))
|
|
|
|
visuailzer = VISUALIZERS.build(dict(type='Visualizer'))
|
|
local_writer = LocalWriter('temp_dir', visuailzer)
|
|
local_writer.add_image('img', image)
|
|
assert os.path.exists(
|
|
os.path.join(local_writer._img_save_dir, 'img_0.png'))
|
|
|
|
shutil.rmtree('temp_dir')
|
|
|
|
def test_add_scalar(self):
|
|
local_writer = LocalWriter('temp_dir')
|
|
local_writer.add_scalar('map', 0.9)
|
|
out_dict = load(local_writer._scalar_save_file, 'json')
|
|
assert out_dict == {'map': 0.9, 'step': 0}
|
|
shutil.rmtree('temp_dir')
|
|
|
|
# test append mode
|
|
local_writer = LocalWriter('temp_dir')
|
|
local_writer.add_scalar('map', 0.9, step=0)
|
|
local_writer.add_scalar('map', 0.95, step=1)
|
|
with open(local_writer._scalar_save_file) as f:
|
|
out_dict = f.read()
|
|
assert out_dict == '{"map": 0.9, "step": 0}\n{"map": ' \
|
|
'0.95, "step": 1}\n'
|
|
shutil.rmtree('temp_dir')
|
|
|
|
def test_add_scalars(self):
|
|
local_writer = LocalWriter('temp_dir')
|
|
input_dict = {'map': 0.7, 'acc': 0.9}
|
|
local_writer.add_scalars(input_dict)
|
|
out_dict = load(local_writer._scalar_save_file, 'json')
|
|
assert out_dict == {'map': 0.7, 'acc': 0.9, 'step': 0}
|
|
|
|
# test append mode
|
|
local_writer.add_scalars({'map': 0.8, 'acc': 0.8}, step=1)
|
|
with open(local_writer._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_writer = LocalWriter('temp_dir')
|
|
local_writer.add_scalars(input_dict, file_path='temp.json')
|
|
assert os.path.exists(local_writer._scalar_save_file)
|
|
assert os.path.exists(
|
|
os.path.join(local_writer._save_dir, 'temp.json'))
|
|
|
|
# file_path and scalar_save_file cannot be the same
|
|
with pytest.raises(AssertionError):
|
|
local_writer.add_scalars(input_dict, file_path='scalars.json')
|
|
|
|
shutil.rmtree('temp_dir')
|
|
|
|
|
|
class TestTensorboardWriter:
|
|
sys.modules['torch.utils.tensorboard'] = MagicMock()
|
|
sys.modules['tensorboardX'] = MagicMock()
|
|
|
|
def test_init(self):
|
|
# visuailzer must be a dictionary or an instance
|
|
# of Visualizer and its subclasses
|
|
with pytest.raises(AssertionError):
|
|
LocalWriter('temp_dir', [dict(type='Visualizer')])
|
|
|
|
TensorboardWriter('temp_dir')
|
|
WRITERS.build(
|
|
dict(
|
|
type='TensorboardWriter',
|
|
visualizer=dict(type='Visualizer'),
|
|
save_dir='temp_dir'))
|
|
|
|
def test_experiment(self):
|
|
tensorboard_writer = TensorboardWriter('temp_dir')
|
|
assert tensorboard_writer.experiment == tensorboard_writer._tensorboard
|
|
|
|
def test_add_graph(self):
|
|
|
|
class Model(nn.Module):
|
|
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.conv = nn.Conv2d(1, 2, 1)
|
|
|
|
def forward(self, x, y=None):
|
|
return self.conv(x)
|
|
|
|
tensorboard_writer = TensorboardWriter('temp_dir')
|
|
|
|
# input must be tensor
|
|
with pytest.raises(AssertionError):
|
|
tensorboard_writer.add_graph(Model(), np.zeros([1, 1, 3, 3]))
|
|
|
|
# input must be 4d tensor
|
|
with pytest.raises(AssertionError):
|
|
tensorboard_writer.add_graph(Model(), torch.zeros([1, 3, 3]))
|
|
|
|
# If the input is a list, the inner element must be a 4d tensor
|
|
with pytest.raises(AssertionError):
|
|
tensorboard_writer.add_graph(
|
|
Model(), [torch.zeros([1, 1, 3, 3]),
|
|
torch.zeros([1, 3, 3])])
|
|
|
|
tensorboard_writer.add_graph(Model(), torch.zeros([1, 1, 3, 3]))
|
|
tensorboard_writer.add_graph(
|
|
Model(), [torch.zeros([1, 1, 3, 3]),
|
|
torch.zeros([1, 1, 3, 3])])
|
|
|
|
def test_add_params(self):
|
|
tensorboard_writer = TensorboardWriter('temp_dir')
|
|
|
|
# 'params_dict' must be dict
|
|
with pytest.raises(AssertionError):
|
|
tensorboard_writer.add_params(['lr', 0])
|
|
|
|
params_dict = dict(lr=0.1, wd=0.2, mode='linear')
|
|
tensorboard_writer.add_params(params_dict)
|
|
|
|
@patch('mmengine.visualization.visualizer.Visualizer.draw', draw)
|
|
def test_add_image(self):
|
|
image = np.random.randint(0, 256, size=(10, 10, 3)).astype(np.uint8)
|
|
|
|
# The visuailzer parameter must be set when
|
|
# the local_writer object is instantiated and
|
|
# the `add_image` method is called.
|
|
with pytest.raises(AssertionError):
|
|
tensorboard_writer = TensorboardWriter('temp_dir')
|
|
tensorboard_writer.add_image('img', image)
|
|
|
|
tensorboard_writer = TensorboardWriter('temp_dir',
|
|
dict(type='Visualizer'))
|
|
tensorboard_writer.add_image('img', image)
|
|
|
|
bboxes = np.array([[1, 1, 2, 2], [1, 1.5, 1, 2.5]])
|
|
tensorboard_writer.visualizer.draw_bboxes(bboxes)
|
|
tensorboard_writer.add_image(
|
|
'img', tensorboard_writer.visualizer.get_image(), step=2)
|
|
|
|
visuailzer = VISUALIZERS.build(dict(type='Visualizer'))
|
|
tensorboard_writer = TensorboardWriter('temp_dir', visuailzer)
|
|
tensorboard_writer.add_image('img', image)
|
|
|
|
def test_add_scalar(self):
|
|
tensorboard_writer = TensorboardWriter('temp_dir')
|
|
tensorboard_writer.add_scalar('map', 0.9)
|
|
# test append mode
|
|
tensorboard_writer.add_scalar('map', 0.9, step=0)
|
|
tensorboard_writer.add_scalar('map', 0.95, step=1)
|
|
|
|
def test_add_scalars(self):
|
|
tensorboard_writer = TensorboardWriter('temp_dir')
|
|
# The step value must be passed through the parameter
|
|
with pytest.raises(AssertionError):
|
|
tensorboard_writer.add_scalars({'map': 0.7, 'acc': 0.9, 'step': 1})
|
|
|
|
input_dict = {'map': 0.7, 'acc': 0.9}
|
|
tensorboard_writer.add_scalars(input_dict)
|
|
# test append mode
|
|
tensorboard_writer.add_scalars({'map': 0.8, 'acc': 0.8}, step=1)
|
|
|
|
|
|
class TestWandbWriter:
|
|
sys.modules['wandb'] = MagicMock()
|
|
|
|
def test_init(self):
|
|
WandbWriter()
|
|
WRITERS.build(
|
|
dict(
|
|
type='WandbWriter',
|
|
visualizer=dict(type='Visualizer'),
|
|
save_dir='temp_dir'))
|
|
|
|
def test_experiment(self):
|
|
wandb_writer = WandbWriter()
|
|
assert wandb_writer.experiment == wandb_writer._wandb
|
|
|
|
def test_add_params(self):
|
|
wandb_writer = WandbWriter()
|
|
|
|
# 'params_dict' must be dict
|
|
with pytest.raises(AssertionError):
|
|
wandb_writer.add_params(['lr', 0])
|
|
|
|
params_dict = dict(lr=0.1, wd=0.2, mode='linear')
|
|
wandb_writer.add_params(params_dict)
|
|
|
|
@patch('mmengine.visualization.visualizer.Visualizer.draw', draw)
|
|
@patch('mmengine.visualization.writer.WandbWriter.add_image_to_wandb',
|
|
Mock)
|
|
def test_add_image(self):
|
|
image = np.random.randint(0, 256, size=(10, 10, 3)).astype(np.uint8)
|
|
|
|
wandb_writer = WandbWriter()
|
|
wandb_writer.add_image('img', image)
|
|
|
|
wandb_writer = WandbWriter(visualizer=dict(type='Visualizer'))
|
|
bboxes = np.array([[1, 1, 2, 2], [1, 1.5, 1, 2.5]])
|
|
wandb_writer.visualizer.set_image(image)
|
|
wandb_writer.visualizer.draw_bboxes(bboxes)
|
|
wandb_writer.add_image(
|
|
'img', wandb_writer.visualizer.get_image(), step=2)
|
|
|
|
visuailzer = VISUALIZERS.build(dict(type='Visualizer'))
|
|
wandb_writer = WandbWriter(visualizer=visuailzer)
|
|
wandb_writer.add_image('img', image)
|
|
|
|
def test_add_scalar(self):
|
|
wandb_writer = WandbWriter()
|
|
wandb_writer.add_scalar('map', 0.9)
|
|
# test append mode
|
|
wandb_writer.add_scalar('map', 0.9, step=0)
|
|
wandb_writer.add_scalar('map', 0.95, step=1)
|
|
|
|
def test_add_scalars(self):
|
|
wandb_writer = WandbWriter()
|
|
input_dict = {'map': 0.7, 'acc': 0.9}
|
|
wandb_writer.add_scalars(input_dict)
|
|
# test append mode
|
|
wandb_writer.add_scalars({'map': 0.8, 'acc': 0.8}, step=1)
|
|
|
|
|
|
class TestComposedWriter:
|
|
sys.modules['torch.utils.tensorboard'] = MagicMock()
|
|
sys.modules['tensorboardX'] = MagicMock()
|
|
sys.modules['wandb'] = MagicMock()
|
|
|
|
def test_init(self):
|
|
|
|
class A:
|
|
pass
|
|
|
|
# The writers inner element must be a dictionary or a
|
|
# subclass of Writer.
|
|
with pytest.raises(AssertionError):
|
|
ComposedWriter(writers=[A()])
|
|
|
|
composed_writer = ComposedWriter(writers=[
|
|
WandbWriter(),
|
|
dict(
|
|
type='TensorboardWriter',
|
|
visualizer=dict(type='Visualizer'),
|
|
save_dir='temp_dir')
|
|
])
|
|
assert len(composed_writer._writers) == 2
|
|
|
|
# test global
|
|
composed_writer = ComposedWriter.create_instance(
|
|
'composed_writer',
|
|
writers=[
|
|
WandbWriter(),
|
|
dict(
|
|
type='TensorboardWriter',
|
|
visualizer=dict(type='Visualizer'),
|
|
save_dir='temp_dir')
|
|
])
|
|
assert len(composed_writer._writers) == 2
|
|
composed_writer_any = ComposedWriter.get_instance('composed_writer')
|
|
assert composed_writer_any == composed_writer
|
|
|
|
def test_get_writer(self):
|
|
composed_writer = ComposedWriter(writers=[
|
|
WandbWriter(),
|
|
dict(
|
|
type='TensorboardWriter',
|
|
visualizer=dict(type='Visualizer'),
|
|
save_dir='temp_dir')
|
|
])
|
|
assert isinstance(composed_writer.get_writer(0), WandbWriter)
|
|
assert isinstance(composed_writer.get_writer(1), TensorboardWriter)
|
|
|
|
def test_get_experiment(self):
|
|
composed_writer = ComposedWriter(writers=[
|
|
WandbWriter(),
|
|
dict(
|
|
type='TensorboardWriter',
|
|
visualizer=dict(type='Visualizer'),
|
|
save_dir='temp_dir')
|
|
])
|
|
assert composed_writer.get_experiment(
|
|
0) == composed_writer._writers[0].experiment
|
|
assert composed_writer.get_experiment(
|
|
1) == composed_writer._writers[1].experiment
|
|
|
|
def test_get_visualizer(self):
|
|
composed_writer = ComposedWriter(writers=[
|
|
WandbWriter(),
|
|
dict(
|
|
type='TensorboardWriter',
|
|
visualizer=dict(type='Visualizer'),
|
|
save_dir='temp_dir')
|
|
])
|
|
assert composed_writer.get_visualizer(
|
|
0) == composed_writer._writers[0].visualizer
|
|
assert composed_writer.get_visualizer(
|
|
1) == composed_writer._writers[1].visualizer
|
|
|
|
def test_add_params(self):
|
|
composed_writer = ComposedWriter(writers=[
|
|
WandbWriter(),
|
|
dict(
|
|
type='TensorboardWriter',
|
|
visualizer=dict(type='Visualizer'),
|
|
save_dir='temp_dir')
|
|
])
|
|
|
|
# 'params_dict' must be dict
|
|
with pytest.raises(AssertionError):
|
|
composed_writer.add_params(['lr', 0])
|
|
|
|
params_dict = dict(lr=0.1, wd=0.2, mode='linear')
|
|
composed_writer.add_params(params_dict)
|
|
|
|
def test_add_graph(self):
|
|
composed_writer = ComposedWriter(writers=[
|
|
WandbWriter(),
|
|
dict(
|
|
type='TensorboardWriter',
|
|
visualizer=dict(type='Visualizer'),
|
|
save_dir='temp_dir')
|
|
])
|
|
|
|
class Model(nn.Module):
|
|
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.conv = nn.Conv2d(1, 2, 1)
|
|
|
|
def forward(self, x, y=None):
|
|
return self.conv(x)
|
|
|
|
# input must be tensor
|
|
with pytest.raises(AssertionError):
|
|
composed_writer.add_graph(Model(), np.zeros([1, 1, 3, 3]))
|
|
|
|
# input must be 4d tensor
|
|
with pytest.raises(AssertionError):
|
|
composed_writer.add_graph(Model(), torch.zeros([1, 3, 3]))
|
|
|
|
# If the input is a list, the inner element must be a 4d tensor
|
|
with pytest.raises(AssertionError):
|
|
composed_writer.add_graph(
|
|
Model(), [torch.zeros([1, 1, 3, 3]),
|
|
torch.zeros([1, 3, 3])])
|
|
|
|
composed_writer.add_graph(Model(), torch.zeros([1, 1, 3, 3]))
|
|
composed_writer.add_graph(
|
|
Model(), [torch.zeros([1, 1, 3, 3]),
|
|
torch.zeros([1, 1, 3, 3])])
|
|
|
|
@patch('mmengine.visualization.visualizer.Visualizer.draw', draw)
|
|
@patch('mmengine.visualization.writer.WandbWriter.add_image_to_wandb',
|
|
Mock)
|
|
def test_add_image(self):
|
|
composed_writer = ComposedWriter(writers=[
|
|
WandbWriter(),
|
|
dict(
|
|
type='TensorboardWriter',
|
|
visualizer=dict(type='Visualizer'),
|
|
save_dir='temp_dir')
|
|
])
|
|
|
|
image = np.random.randint(0, 256, size=(10, 10, 3)).astype(np.uint8)
|
|
composed_writer.add_image('img', image)
|
|
|
|
bboxes = np.array([[1, 1, 2, 2], [1, 1.5, 1, 2.5]])
|
|
composed_writer.get_writer(1).visualizer.draw_bboxes(bboxes)
|
|
composed_writer.get_writer(1).add_image(
|
|
'img',
|
|
composed_writer.get_writer(1).visualizer.get_image(),
|
|
step=2)
|
|
|
|
def test_add_scalar(self):
|
|
composed_writer = ComposedWriter(writers=[
|
|
WandbWriter(),
|
|
dict(
|
|
type='TensorboardWriter',
|
|
visualizer=dict(type='Visualizer'),
|
|
save_dir='temp_dir')
|
|
])
|
|
composed_writer.add_scalar('map', 0.9)
|
|
# test append mode
|
|
composed_writer.add_scalar('map', 0.9, step=0)
|
|
composed_writer.add_scalar('map', 0.95, step=1)
|
|
|
|
def test_add_scalars(self):
|
|
composed_writer = ComposedWriter(writers=[
|
|
WandbWriter(),
|
|
dict(
|
|
type='TensorboardWriter',
|
|
visualizer=dict(type='Visualizer'),
|
|
save_dir='temp_dir')
|
|
])
|
|
input_dict = {'map': 0.7, 'acc': 0.9}
|
|
composed_writer.add_scalars(input_dict)
|
|
# test append mode
|
|
composed_writer.add_scalars({'map': 0.8, 'acc': 0.8}, step=1)
|