# -*- coding: utf-8 -*-
# Copyright (c) Alibaba, Inc. and its affiliates.
import copy
import os.path
from numbers import Number
import numpy as np
import torch
from PIL import Image
from scepter.modules.utils.file_system import FS
def check_legal_type(data):
if isinstance(data, str) or isinstance(data, Number):
return True
elif isinstance(data, dict):
for k, v in data.items():
if not check_legal_type(v):
return False
return True
elif isinstance(data, list):
for v in data:
if not check_legal_type(v):
return False
return True
else:
return False
def register_data(probe_data: dict, key_prefix=''):
ret_data = {}
dist_data = {}
for k, v in probe_data.items():
key = f'{key_prefix}_{k}'
if isinstance(v, torch.Tensor) or isinstance(v, np.ndarray):
ret_data[key] = ProbeData(v)
elif isinstance(v, ProbeData):
ret_data[key] = v
else:
if not check_legal_type(v):
raise f'The datatype of {key} should be included in [array, tensor, number, str] or the dict or ' \
f'list of (number, str); if you want register the list of image, please use ProbeData instance.'
ret_data[key] = ProbeData(v)
if ret_data[key].view_distribute:
dist_data[key] = ret_data[key].distribute
return ret_data, dist_data
def merge_gathered_probe(all_gathered_data):
'''
Merge the gathered data on rank_0.
Returns:
The merged data.
'''
for key, gathered_data in all_gathered_data.items():
# Must be the list of ProbeData.
if isinstance(gathered_data, list):
for v in gathered_data:
if not isinstance(v, ProbeData):
all_gathered_data[key] = gathered_data
# Must be the gathered data.
ret_data = gathered_data[0]
if not isinstance(ret_data.data,
list) and (isinstance(ret_data.data, np.ndarray)
or isinstance(ret_data.data, dict)
or check_legal_type(ret_data.data)):
new_data = [v.data for v in gathered_data]
if ret_data.build_label is not None:
ret_data.build_label = [
v.build_label for v in gathered_data
]
all_gathered_data[key] = ProbeData(
new_data,
is_image=ret_data.is_image,
build_html=ret_data.build_html,
build_label=ret_data.build_label,
view_distribute=ret_data.view_distribute)
elif isinstance(ret_data.data, list):
if ret_data.build_label is not None:
if isinstance(ret_data.build_label, str):
ret_data.build_label = [
ret_data.build_label for _ in ret_data.data
]
for v in gathered_data[1:]:
ret_data.data += v.data
if ret_data.build_label is not None:
if isinstance(v.build_label, str):
ret_data.build_label.extend(
[v.build_label for _ in v.data])
ret_data.build_label.extend(v.build_label)
all_gathered_data[key] = ProbeData(
ret_data.data,
is_image=ret_data.is_image,
build_html=ret_data.build_html,
build_label=ret_data.build_label,
view_distribute=ret_data.view_distribute)
else:
all_gathered_data[key] = gathered_data
return all_gathered_data
class ProbeData():
def __init__(self,
data,
is_image=False,
build_html=False,
build_label=None,
view_distribute=False):
''' Probe Data Initialize.
We only support basic types such as [torch.Tensor, numpy.ndarray, number, str],
or [dict, list] of [dict, list,
number, str], or [dict, list] of [tensor, array]
'''
data = copy.deepcopy(data)
self.basic_type = True
self._distribute_dict = {}
if view_distribute:
is_legal = True
if isinstance(data, str) or isinstance(data, Number):
is_legal = True
if data in self._distribute_dict:
self._distribute_dict[data] += 1
else:
self._distribute_dict[data] = 1
elif isinstance(data, list):
for v in data:
if isinstance(v, str) or isinstance(v, Number):
is_legal = True
if v in self._distribute_dict:
self._distribute_dict[v] += 1
else:
self._distribute_dict[v] = 1
else:
is_legal = False
elif isinstance(data, dict):
for k, v in data.items():
if isinstance(v, str) or isinstance(v, Number):
is_legal = True
n_k = f'{k}_{v}'
if n_k in self._distribute_dict:
self._distribute_dict[n_k] += 1
else:
self._distribute_dict[n_k] = 1
else:
is_legal = False
else:
is_legal = False
if not is_legal:
print('Unsurpport data type', data)
assert is_legal
self.view_distribute = view_distribute
if isinstance(data, torch.Tensor):
self.data = data.detach().cpu().numpy()
elif isinstance(data, np.ndarray):
self.data = data
elif isinstance(data, dict):
for k, v in data.items():
if not check_legal_type(v):
if isinstance(v, torch.Tensor):
data[k] = v.detach().cpu().numpy()
self.basic_type = False
elif isinstance(v, np.ndarray):
data[k] = v
self.basic_type = False
else:
raise f'Unsupport data type for {v}'
self.data = data
elif isinstance(data, list):
for idx, v in enumerate(data):
if not check_legal_type(v):
if isinstance(v, torch.Tensor):
data[idx] = v.detach().cpu().numpy()
self.basic_type = False
elif isinstance(v, np.ndarray):
data[idx] = v
self.basic_type = False
else:
raise f'Unsupport data type for {v}'
self.data = data
elif check_legal_type(data):
self.data = data
else:
raise f'Unsupport data type for {data}'
self.is_image = is_image
self.build_html = build_html
if self.build_html:
assert build_label is not None
if isinstance(self.data, str):
assert isinstance(build_label, str)
if isinstance(self.data, list):
assert isinstance(build_label, str) or isinstance(
build_label, list)
if isinstance(self.data, dict):
assert isinstance(build_label, str) or isinstance(
build_label, dict)
self.build_label = build_label
def save_image(self, file_prefix, images):
np_shape = images.shape
# 4D
shape_str = '_'.join([str(v) for v in np_shape])
if len(np_shape) == 4:
# channel is 1 or 3
if np_shape[-1] == 1 or np_shape[-1] == 3:
if np_shape[-1] == 1:
images = images.reshape(images.shape[:-1])
file_list = []
for idx in range(np_shape[0]):
file_path = file_prefix + f'_probe_{idx}_[{shape_str}].jpg'
with FS.put_to(file_path) as local_path:
Image.fromarray(images[idx, ...]).save(local_path)
file_list.append(file_path)
return file_list
else:
raise f"Ensure your data's dim is BWHC, and channel is 1 or 3 for {file_prefix}"
elif len(np_shape) == 3:
if np_shape[-1] == 1 or np_shape[-1] == 3:
if np_shape[-1] == 1:
images = images.reshape(images.shape[:-1])
file_path = file_prefix + f'_probe_[{shape_str}].jpg'
with FS.put_to(file_path) as local_path:
Image.fromarray(images).save(local_path)
return file_path
else:
images = images.reshape(list(images.shape) + [1])
return self.save_image(file_prefix, images)
elif len(np_shape) == 2:
file_path = file_prefix + f'_probe_[{shape_str}].jpg'
with FS.put_to(file_path) as local_path:
Image.fromarray(images).save(local_path)
return file_path
else:
raise f"Ensure your data's dim is BWHC or WHC or WH, and channel is 1 or 3 for {file_prefix}"
def save_npy(self, file_prefix, data):
shape_str = '_'.join([str(v) for v in data.shape])
file_path = file_prefix + f'_{shape_str}.npy'
with FS.put_to(file_path) as local_path:
np.save(local_path, data)
return file_path
def save_html(self, html_prefix, ret_data, ret_label):
height = 600
with FS.put_to(html_prefix) as local_path:
with open(local_path, 'w') as f:
f.writelines('\n')
f.writelines('\n')
f.writelines(' \n')
all_ranks = list()
for save_id, save_data in enumerate(zip(ret_data, ret_label)):
save_path, save_label = save_data
one_rank = '