rename lib to opengait
This commit is contained in:
darkliang
@@ -0,0 +1,10 @@
|
||||
from .common import get_ddp_module, ddp_all_gather
|
||||
from .common import Odict, Ntuple
|
||||
from .common import get_valid_args
|
||||
from .common import is_list_or_tuple, is_bool, is_str, is_list, is_dict, is_tensor, is_array, config_loader, init_seeds, handler, params_count
|
||||
from .common import ts2np, ts2var, np2var, list2var
|
||||
from .common import mkdir, clones
|
||||
from .common import MergeCfgsDict
|
||||
from .common import get_attr_from
|
||||
from .common import NoOp
|
||||
from .msg_manager import get_msg_mgr
|
||||
@@ -0,0 +1,205 @@
|
||||
import copy
|
||||
import os
|
||||
import inspect
|
||||
import logging
|
||||
import torch
|
||||
import numpy as np
|
||||
import torch.nn as nn
|
||||
import torch.autograd as autograd
|
||||
import yaml
|
||||
import random
|
||||
from torch.nn.parallel import DistributedDataParallel as DDP
|
||||
from collections import OrderedDict, namedtuple
|
||||
|
||||
|
||||
class NoOp:
|
||||
def __getattr__(self, *args):
|
||||
def no_op(*args, **kwargs): pass
|
||||
return no_op
|
||||
|
||||
|
||||
class Odict(OrderedDict):
|
||||
def append(self, odict):
|
||||
dst_keys = self.keys()
|
||||
for k, v in odict.items():
|
||||
if not is_list(v):
|
||||
v = [v]
|
||||
if k in dst_keys:
|
||||
if is_list(self[k]):
|
||||
self[k] += v
|
||||
else:
|
||||
self[k] = [self[k]] + v
|
||||
else:
|
||||
self[k] = v
|
||||
|
||||
|
||||
def Ntuple(description, keys, values):
|
||||
if not is_list_or_tuple(keys):
|
||||
keys = [keys]
|
||||
values = [values]
|
||||
Tuple = namedtuple(description, keys)
|
||||
return Tuple._make(values)
|
||||
|
||||
|
||||
def get_valid_args(obj, input_args, free_keys=[]):
|
||||
if inspect.isfunction(obj):
|
||||
expected_keys = inspect.getargspec(obj)[0]
|
||||
elif inspect.isclass(obj):
|
||||
expected_keys = inspect.getargspec(obj.__init__)[0]
|
||||
else:
|
||||
raise ValueError('Just support function and class object!')
|
||||
unexpect_keys = list()
|
||||
expected_args = {}
|
||||
for k, v in input_args.items():
|
||||
if k in expected_keys:
|
||||
expected_args[k] = v
|
||||
elif k in free_keys:
|
||||
pass
|
||||
else:
|
||||
unexpect_keys.append(k)
|
||||
if unexpect_keys != []:
|
||||
logging.info("Find Unexpected Args(%s) in the Configuration of - %s -" %
|
||||
(', '.join(unexpect_keys), obj.__name__))
|
||||
return expected_args
|
||||
|
||||
|
||||
def get_attr_from(sources, name):
|
||||
try:
|
||||
return getattr(sources[0], name)
|
||||
except:
|
||||
return get_attr_from(sources[1:], name) if len(sources) > 1 else getattr(sources[0], name)
|
||||
|
||||
|
||||
def is_list_or_tuple(x):
|
||||
return isinstance(x, (list, tuple))
|
||||
|
||||
|
||||
def is_bool(x):
|
||||
return isinstance(x, bool)
|
||||
|
||||
|
||||
def is_str(x):
|
||||
return isinstance(x, str)
|
||||
|
||||
|
||||
def is_list(x):
|
||||
return isinstance(x, list) or isinstance(x, nn.ModuleList)
|
||||
|
||||
|
||||
def is_dict(x):
|
||||
return isinstance(x, dict) or isinstance(x, OrderedDict) or isinstance(x, Odict)
|
||||
|
||||
|
||||
def is_tensor(x):
|
||||
return isinstance(x, torch.Tensor)
|
||||
|
||||
|
||||
def is_array(x):
|
||||
return isinstance(x, np.ndarray)
|
||||
|
||||
|
||||
def ts2np(x):
|
||||
return x.cpu().data.numpy()
|
||||
|
||||
|
||||
def ts2var(x, **kwargs):
|
||||
return autograd.Variable(x, **kwargs).cuda()
|
||||
|
||||
|
||||
def np2var(x, **kwargs):
|
||||
return ts2var(torch.from_numpy(x), **kwargs)
|
||||
|
||||
|
||||
def list2var(x, **kwargs):
|
||||
return np2var(np.array(x), **kwargs)
|
||||
|
||||
|
||||
def mkdir(path):
|
||||
if not os.path.exists(path):
|
||||
os.makedirs(path)
|
||||
|
||||
|
||||
def MergeCfgsDict(src, dst):
|
||||
for k, v in src.items():
|
||||
if (k not in dst.keys()) or (type(v) != type(dict())):
|
||||
dst[k] = v
|
||||
else:
|
||||
if is_dict(src[k]) and is_dict(dst[k]):
|
||||
MergeCfgsDict(src[k], dst[k])
|
||||
else:
|
||||
dst[k] = v
|
||||
|
||||
|
||||
def clones(module, N):
|
||||
"Produce N identical layers."
|
||||
return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])
|
||||
|
||||
|
||||
def config_loader(path):
|
||||
with open(path, 'r') as stream:
|
||||
src_cfgs = yaml.safe_load(stream)
|
||||
with open("./config/default.yaml", 'r') as stream:
|
||||
dst_cfgs = yaml.safe_load(stream)
|
||||
MergeCfgsDict(src_cfgs, dst_cfgs)
|
||||
return dst_cfgs
|
||||
|
||||
|
||||
def init_seeds(seed=0, cuda_deterministic=True):
|
||||
random.seed(seed)
|
||||
np.random.seed(seed)
|
||||
torch.manual_seed(seed)
|
||||
torch.cuda.manual_seed_all(seed)
|
||||
# Speed-reproducibility tradeoff https://pytorch.org/docs/stable/notes/randomness.html
|
||||
if cuda_deterministic: # slower, more reproducible
|
||||
torch.backends.cudnn.deterministic = True
|
||||
torch.backends.cudnn.benchmark = False
|
||||
else: # faster, less reproducible
|
||||
torch.backends.cudnn.deterministic = False
|
||||
torch.backends.cudnn.benchmark = True
|
||||
|
||||
|
||||
def handler(signum, frame):
|
||||
logging.info('Ctrl+c/z pressed')
|
||||
os.system(
|
||||
"kill $(ps aux | grep main.py | grep -v grep | awk '{print $2}') ")
|
||||
logging.info('process group flush!')
|
||||
|
||||
|
||||
def ddp_all_gather(features, dim=0, requires_grad=True):
|
||||
'''
|
||||
inputs: [n, ...]
|
||||
'''
|
||||
|
||||
world_size = torch.distributed.get_world_size()
|
||||
rank = torch.distributed.get_rank()
|
||||
feature_list = [torch.ones_like(features) for _ in range(world_size)]
|
||||
torch.distributed.all_gather(feature_list, features.contiguous())
|
||||
|
||||
if requires_grad:
|
||||
feature_list[rank] = features
|
||||
feature = torch.cat(feature_list, dim=dim)
|
||||
return feature
|
||||
|
||||
|
||||
# https://github.com/pytorch/pytorch/issues/16885
|
||||
class DDPPassthrough(DDP):
|
||||
def __getattr__(self, name):
|
||||
try:
|
||||
return super().__getattr__(name)
|
||||
except AttributeError:
|
||||
return getattr(self.module, name)
|
||||
|
||||
|
||||
def get_ddp_module(module, **kwargs):
|
||||
if len(list(module.parameters())) == 0:
|
||||
# for the case that loss module has not parameters.
|
||||
return module
|
||||
device = torch.cuda.current_device()
|
||||
module = DDPPassthrough(module, device_ids=[device], output_device=device,
|
||||
find_unused_parameters=False, **kwargs)
|
||||
return module
|
||||
|
||||
|
||||
def params_count(net):
|
||||
n_parameters = sum(p.numel() for p in net.parameters())
|
||||
return 'Parameters Count: {:.5f}M'.format(n_parameters / 1e6)
|
||||
@@ -0,0 +1,276 @@
|
||||
import os
|
||||
from time import strftime, localtime
|
||||
import torch
|
||||
import numpy as np
|
||||
import torch.nn.functional as F
|
||||
from utils import get_msg_mgr, mkdir
|
||||
|
||||
|
||||
def cuda_dist(x, y, metric='euc'):
|
||||
x = torch.from_numpy(x).cuda()
|
||||
y = torch.from_numpy(y).cuda()
|
||||
if metric == 'cos':
|
||||
x = F.normalize(x, p=2, dim=2) # n p c
|
||||
y = F.normalize(y, p=2, dim=2) # n p c
|
||||
num_bin = x.size(1)
|
||||
n_x = x.size(0)
|
||||
n_y = y.size(0)
|
||||
dist = torch.zeros(n_x, n_y).cuda()
|
||||
for i in range(num_bin):
|
||||
_x = x[:, i, ...]
|
||||
_y = y[:, i, ...]
|
||||
if metric == 'cos':
|
||||
dist += torch.matmul(_x, _y.transpose(0, 1))
|
||||
else:
|
||||
_dist = torch.sum(_x ** 2, 1).unsqueeze(1) + torch.sum(_y ** 2, 1).unsqueeze(
|
||||
1).transpose(0, 1) - 2 * torch.matmul(_x, _y.transpose(0, 1))
|
||||
dist += torch.sqrt(F.relu(_dist))
|
||||
return 1 - dist/num_bin if metric == 'cos' else dist / num_bin
|
||||
|
||||
# Exclude identical-view cases
|
||||
|
||||
|
||||
def de_diag(acc, each_angle=False):
|
||||
dividend = acc.shape[1] - 1.
|
||||
result = np.sum(acc - np.diag(np.diag(acc)), 1) / dividend
|
||||
if not each_angle:
|
||||
result = np.mean(result)
|
||||
return result
|
||||
|
||||
# Modified From https://github.com/AbnerHqC/GaitSet/blob/master/model/utils/evaluator.py
|
||||
|
||||
|
||||
def identification(data, dataset, metric='euc'):
|
||||
msg_mgr = get_msg_mgr()
|
||||
feature, label, seq_type, view = data['embeddings'], data['labels'], data['types'], data['views']
|
||||
label = np.array(label)
|
||||
view_list = list(set(view))
|
||||
view_list.sort()
|
||||
view_num = len(view_list)
|
||||
# sample_num = len(feature)
|
||||
|
||||
probe_seq_dict = {'CASIA-B': [['nm-05', 'nm-06'], ['bg-01', 'bg-02'], ['cl-01', 'cl-02']],
|
||||
'OUMVLP': [['00']]}
|
||||
|
||||
gallery_seq_dict = {'CASIA-B': [['nm-01', 'nm-02', 'nm-03', 'nm-04']],
|
||||
'OUMVLP': [['01']]}
|
||||
if dataset not in (probe_seq_dict or gallery_seq_dict):
|
||||
raise KeyError("DataSet %s hasn't been supported !" % dataset)
|
||||
num_rank = 5
|
||||
acc = np.zeros([len(probe_seq_dict[dataset]),
|
||||
view_num, view_num, num_rank]) - 1.
|
||||
for (p, probe_seq) in enumerate(probe_seq_dict[dataset]):
|
||||
for gallery_seq in gallery_seq_dict[dataset]:
|
||||
for (v1, probe_view) in enumerate(view_list):
|
||||
for (v2, gallery_view) in enumerate(view_list):
|
||||
gseq_mask = np.isin(seq_type, gallery_seq) & np.isin(
|
||||
view, [gallery_view])
|
||||
gallery_x = feature[gseq_mask, :]
|
||||
gallery_y = label[gseq_mask]
|
||||
|
||||
pseq_mask = np.isin(seq_type, probe_seq) & np.isin(
|
||||
view, [probe_view])
|
||||
probe_x = feature[pseq_mask, :]
|
||||
probe_y = label[pseq_mask]
|
||||
|
||||
dist = cuda_dist(probe_x, gallery_x, metric)
|
||||
idx = dist.sort(1)[1].cpu().numpy()
|
||||
acc[p, v1, v2, :] = np.round(
|
||||
np.sum(np.cumsum(np.reshape(probe_y, [-1, 1]) == gallery_y[idx[:, 0:num_rank]], 1) > 0,
|
||||
0) * 100 / dist.shape[0], 2)
|
||||
result_dict = {}
|
||||
np.set_printoptions(precision=3, suppress=True)
|
||||
if 'OUMVLP' not in dataset:
|
||||
for i in range(1):
|
||||
msg_mgr.log_info(
|
||||
'===Rank-%d (Include identical-view cases)===' % (i + 1))
|
||||
msg_mgr.log_info('NM: %.3f,\tBG: %.3f,\tCL: %.3f' % (
|
||||
np.mean(acc[0, :, :, i]),
|
||||
np.mean(acc[1, :, :, i]),
|
||||
np.mean(acc[2, :, :, i])))
|
||||
for i in range(1):
|
||||
msg_mgr.log_info(
|
||||
'===Rank-%d (Exclude identical-view cases)===' % (i + 1))
|
||||
msg_mgr.log_info('NM: %.3f,\tBG: %.3f,\tCL: %.3f' % (
|
||||
de_diag(acc[0, :, :, i]),
|
||||
de_diag(acc[1, :, :, i]),
|
||||
de_diag(acc[2, :, :, i])))
|
||||
result_dict["scalar/test_accuracy/NM"] = de_diag(acc[0, :, :, i])
|
||||
result_dict["scalar/test_accuracy/BG"] = de_diag(acc[1, :, :, i])
|
||||
result_dict["scalar/test_accuracy/CL"] = de_diag(acc[2, :, :, i])
|
||||
np.set_printoptions(precision=2, floatmode='fixed')
|
||||
for i in range(1):
|
||||
msg_mgr.log_info(
|
||||
'===Rank-%d of each angle (Exclude identical-view cases)===' % (i + 1))
|
||||
msg_mgr.log_info('NM: {}'.format(de_diag(acc[0, :, :, i], True)))
|
||||
msg_mgr.log_info('BG: {}'.format(de_diag(acc[1, :, :, i], True)))
|
||||
msg_mgr.log_info('CL: {}'.format(de_diag(acc[2, :, :, i], True)))
|
||||
else:
|
||||
msg_mgr.log_info('===Rank-1 (Include identical-view cases)===')
|
||||
msg_mgr.log_info('NM: %.3f ' % (np.mean(acc[0, :, :, 0])))
|
||||
msg_mgr.log_info('===Rank-1 (Exclude identical-view cases)===')
|
||||
msg_mgr.log_info('NM: %.3f ' % (de_diag(acc[0, :, :, 0])))
|
||||
msg_mgr.log_info(
|
||||
'===Rank-1 of each angle (Exclude identical-view cases)===')
|
||||
msg_mgr.log_info('NM: {}'.format(de_diag(acc[0, :, :, 0], True)))
|
||||
result_dict["scalar/test_accuracy/NM"] = de_diag(acc[0, :, :, 0])
|
||||
return result_dict
|
||||
|
||||
|
||||
def identification_real_scene(data, dataset, metric='euc'):
|
||||
msg_mgr = get_msg_mgr()
|
||||
feature, label, seq_type = data['embeddings'], data['labels'], data['types']
|
||||
label = np.array(label)
|
||||
|
||||
gallery_seq_type = {'0001-1000': ['1', '2'],
|
||||
"HID2021": ['0'], '0001-1000-test': ['0'],
|
||||
'GREW': ['01']}
|
||||
probe_seq_type = {'0001-1000': ['3', '4', '5', '6'],
|
||||
"HID2021": ['1'], '0001-1000-test': ['1'],
|
||||
'GREW': ['02']}
|
||||
|
||||
num_rank = 20
|
||||
acc = np.zeros([num_rank]) - 1.
|
||||
gseq_mask = np.isin(seq_type, gallery_seq_type[dataset])
|
||||
gallery_x = feature[gseq_mask, :]
|
||||
gallery_y = label[gseq_mask]
|
||||
pseq_mask = np.isin(seq_type, probe_seq_type[dataset])
|
||||
probe_x = feature[pseq_mask, :]
|
||||
probe_y = label[pseq_mask]
|
||||
|
||||
dist = cuda_dist(probe_x, gallery_x, metric)
|
||||
idx = dist.cpu().sort(1)[1].numpy()
|
||||
acc = np.round(np.sum(np.cumsum(np.reshape(probe_y, [-1, 1]) == gallery_y[idx[:, 0:num_rank]], 1) > 0,
|
||||
0) * 100 / dist.shape[0], 2)
|
||||
msg_mgr.log_info('==Rank-1==')
|
||||
msg_mgr.log_info('%.3f' % (np.mean(acc[0])))
|
||||
msg_mgr.log_info('==Rank-5==')
|
||||
msg_mgr.log_info('%.3f' % (np.mean(acc[4])))
|
||||
msg_mgr.log_info('==Rank-10==')
|
||||
msg_mgr.log_info('%.3f' % (np.mean(acc[9])))
|
||||
msg_mgr.log_info('==Rank-20==')
|
||||
msg_mgr.log_info('%.3f' % (np.mean(acc[19])))
|
||||
return {"scalar/test_accuracy/Rank-1": np.mean(acc[0]), "scalar/test_accuracy/Rank-5": np.mean(acc[4])}
|
||||
|
||||
|
||||
def identification_GREW_submission(data, dataset, metric='euc'):
|
||||
get_msg_mgr().log_info("Evaluating GREW")
|
||||
feature, label, seq_type, view = data['embeddings'], data['labels'], data['types'], data['views']
|
||||
label = np.array(label)
|
||||
view = np.array(view)
|
||||
gallery_seq_type = {'GREW': ['01', '02']}
|
||||
probe_seq_type = {'GREW': ['03']}
|
||||
gseq_mask = np.isin(seq_type, gallery_seq_type[dataset])
|
||||
gallery_x = feature[gseq_mask, :]
|
||||
gallery_y = label[gseq_mask]
|
||||
pseq_mask = np.isin(seq_type, probe_seq_type[dataset])
|
||||
probe_x = feature[pseq_mask, :]
|
||||
probe_y = view[pseq_mask]
|
||||
|
||||
dist = cuda_dist(probe_x, gallery_x, metric)
|
||||
idx = dist.cpu().sort(1)[1].numpy()
|
||||
|
||||
save_path = os.path.join(
|
||||
"GREW_result/"+strftime('%Y-%m%d-%H%M%S', localtime())+".csv")
|
||||
mkdir("GREW_result")
|
||||
with open(save_path, "w") as f:
|
||||
f.write("videoId,rank1,rank2,rank3,rank4,rank5,rank6,rank7,rank8,rank9,rank10,rank11,rank12,rank13,rank14,rank15,rank16,rank17,rank18,rank19,rank20\n")
|
||||
for i in range(len(idx)):
|
||||
r_format = [int(idx) for idx in gallery_y[idx[i, 0:20]]]
|
||||
output_row = '{}'+',{}'*20+'\n'
|
||||
f.write(output_row.format(probe_y[i], *r_format))
|
||||
print("GREW result saved to {}/{}".format(os.getcwd(), save_path))
|
||||
return
|
||||
|
||||
|
||||
def evaluate_HID(data, dataset, metric='euc'):
|
||||
msg_mgr = get_msg_mgr()
|
||||
msg_mgr.log_info("Evaluating HID")
|
||||
feature, label, seq_type = data['embeddings'], data['labels'], data['types']
|
||||
label = np.array(label)
|
||||
seq_type = np.array(seq_type)
|
||||
probe_mask = (label == "probe")
|
||||
gallery_mask = (label != "probe")
|
||||
gallery_x = feature[gallery_mask, :]
|
||||
gallery_y = label[gallery_mask]
|
||||
probe_x = feature[probe_mask, :]
|
||||
probe_y = seq_type[probe_mask]
|
||||
|
||||
feat = np.concatenate([probe_x, gallery_x])
|
||||
dist = cuda_dist(feat, feat, metric).cpu().numpy()
|
||||
msg_mgr.log_info("Starting Re-ranking")
|
||||
re_rank = re_ranking(dist, probe_x.shape[0], k1=6, k2=6, lambda_value=0.3)
|
||||
idx = np.argsort(re_rank, axis=1)
|
||||
|
||||
save_path = os.path.join(
|
||||
"HID_result/"+strftime('%Y-%m%d-%H%M%S', localtime())+".csv")
|
||||
mkdir("HID_result")
|
||||
with open(save_path, "w") as f:
|
||||
f.write("videoID,label\n")
|
||||
for i in range(len(idx)):
|
||||
f.write("{},{}\n".format(probe_y[i], gallery_y[idx[i, 0]]))
|
||||
print("HID result saved to {}/{}".format(os.getcwd(), save_path))
|
||||
return
|
||||
|
||||
|
||||
def re_ranking(original_dist, query_num, k1, k2, lambda_value):
|
||||
# Modified from https://github.com/michuanhaohao/reid-strong-baseline/blob/master/utils/re_ranking.py
|
||||
all_num = original_dist.shape[0]
|
||||
original_dist = np.transpose(original_dist / np.max(original_dist, axis=0))
|
||||
V = np.zeros_like(original_dist).astype(np.float16)
|
||||
initial_rank = np.argsort(original_dist).astype(np.int32)
|
||||
|
||||
for i in range(all_num):
|
||||
# k-reciprocal neighbors
|
||||
forward_k_neigh_index = initial_rank[i, :k1 + 1]
|
||||
backward_k_neigh_index = initial_rank[forward_k_neigh_index, :k1 + 1]
|
||||
fi = np.where(backward_k_neigh_index == i)[0]
|
||||
k_reciprocal_index = forward_k_neigh_index[fi]
|
||||
k_reciprocal_expansion_index = k_reciprocal_index
|
||||
for j in range(len(k_reciprocal_index)):
|
||||
candidate = k_reciprocal_index[j]
|
||||
candidate_forward_k_neigh_index = initial_rank[candidate, :int(
|
||||
np.around(k1 / 2)) + 1]
|
||||
candidate_backward_k_neigh_index = initial_rank[candidate_forward_k_neigh_index,
|
||||
:int(np.around(k1 / 2)) + 1]
|
||||
fi_candidate = np.where(
|
||||
candidate_backward_k_neigh_index == candidate)[0]
|
||||
candidate_k_reciprocal_index = candidate_forward_k_neigh_index[fi_candidate]
|
||||
if len(np.intersect1d(candidate_k_reciprocal_index, k_reciprocal_index)) > 2 / 3 * len(
|
||||
candidate_k_reciprocal_index):
|
||||
k_reciprocal_expansion_index = np.append(
|
||||
k_reciprocal_expansion_index, candidate_k_reciprocal_index)
|
||||
|
||||
k_reciprocal_expansion_index = np.unique(k_reciprocal_expansion_index)
|
||||
weight = np.exp(-original_dist[i, k_reciprocal_expansion_index])
|
||||
V[i, k_reciprocal_expansion_index] = weight / np.sum(weight)
|
||||
original_dist = original_dist[:query_num, ]
|
||||
if k2 != 1:
|
||||
V_qe = np.zeros_like(V, dtype=np.float16)
|
||||
for i in range(all_num):
|
||||
V_qe[i, :] = np.mean(V[initial_rank[i, :k2], :], axis=0)
|
||||
V = V_qe
|
||||
del V_qe
|
||||
del initial_rank
|
||||
invIndex = []
|
||||
for i in range(all_num):
|
||||
invIndex.append(np.where(V[:, i] != 0)[0])
|
||||
|
||||
jaccard_dist = np.zeros_like(original_dist, dtype=np.float16)
|
||||
|
||||
for i in range(query_num):
|
||||
temp_min = np.zeros(shape=[1, all_num], dtype=np.float16)
|
||||
indNonZero = np.where(V[i, :] != 0)[0]
|
||||
indImages = [invIndex[ind] for ind in indNonZero]
|
||||
for j in range(len(indNonZero)):
|
||||
temp_min[0, indImages[j]] = temp_min[0, indImages[j]] + np.minimum(V[i, indNonZero[j]],
|
||||
V[indImages[j], indNonZero[j]])
|
||||
jaccard_dist[i] = 1 - temp_min / (2 - temp_min)
|
||||
|
||||
final_dist = jaccard_dist * (1 - lambda_value) + \
|
||||
original_dist * lambda_value
|
||||
del original_dist
|
||||
del V
|
||||
del jaccard_dist
|
||||
final_dist = final_dist[:query_num, query_num:]
|
||||
return final_dist
|
||||
@@ -0,0 +1,121 @@
|
||||
import time
|
||||
import torch
|
||||
|
||||
import numpy as np
|
||||
import torchvision.utils as vutils
|
||||
import os.path as osp
|
||||
from time import strftime, localtime
|
||||
|
||||
from torch.utils.tensorboard import SummaryWriter
|
||||
from .common import is_list, is_tensor, ts2np, mkdir, Odict, NoOp
|
||||
import logging
|
||||
|
||||
|
||||
class MessageManager:
|
||||
def __init__(self):
|
||||
self.info_dict = Odict()
|
||||
self.writer_hparams = ['image', 'scalar']
|
||||
self.time = time.time()
|
||||
|
||||
def init_manager(self, save_path, log_to_file, log_iter, iteration=0):
|
||||
self.iteration = iteration
|
||||
self.log_iter = log_iter
|
||||
mkdir(osp.join(save_path, "summary/"))
|
||||
self.writer = SummaryWriter(
|
||||
osp.join(save_path, "summary/"), purge_step=self.iteration)
|
||||
self.init_logger(save_path, log_to_file)
|
||||
|
||||
def init_logger(self, save_path, log_to_file):
|
||||
# init logger
|
||||
self.logger = logging.getLogger('opengait')
|
||||
self.logger.setLevel(logging.INFO)
|
||||
self.logger.propagate = False
|
||||
formatter = logging.Formatter(
|
||||
fmt='[%(asctime)s] [%(levelname)s]: %(message)s', datefmt='%Y-%m-%d %H:%M:%S')
|
||||
if log_to_file:
|
||||
mkdir(osp.join(save_path, "logs/"))
|
||||
vlog = logging.FileHandler(
|
||||
osp.join(save_path, "logs/", strftime('%Y-%m-%d-%H-%M-%S', localtime())+'.txt'))
|
||||
vlog.setLevel(logging.INFO)
|
||||
vlog.setFormatter(formatter)
|
||||
self.logger.addHandler(vlog)
|
||||
|
||||
console = logging.StreamHandler()
|
||||
console.setFormatter(formatter)
|
||||
console.setLevel(logging.DEBUG)
|
||||
self.logger.addHandler(console)
|
||||
|
||||
def append(self, info):
|
||||
for k, v in info.items():
|
||||
v = [v] if not is_list(v) else v
|
||||
v = [ts2np(_) if is_tensor(_) else _ for _ in v]
|
||||
info[k] = v
|
||||
self.info_dict.append(info)
|
||||
|
||||
def flush(self):
|
||||
self.info_dict.clear()
|
||||
self.writer.flush()
|
||||
|
||||
def write_to_tensorboard(self, summary):
|
||||
|
||||
for k, v in summary.items():
|
||||
module_name = k.split('/')[0]
|
||||
if module_name not in self.writer_hparams:
|
||||
self.log_warning(
|
||||
'Not Expected --Summary-- type [{}] appear!!!{}'.format(k, self.writer_hparams))
|
||||
continue
|
||||
board_name = k.replace(module_name + "/", '')
|
||||
writer_module = getattr(self.writer, 'add_' + module_name)
|
||||
v = v.detach() if is_tensor(v) else v
|
||||
v = vutils.make_grid(
|
||||
v, normalize=True, scale_each=True) if 'image' in module_name else v
|
||||
if module_name == 'scalar':
|
||||
try:
|
||||
v = v.mean()
|
||||
except:
|
||||
v = v
|
||||
writer_module(board_name, v, self.iteration)
|
||||
|
||||
def log_training_info(self):
|
||||
now = time.time()
|
||||
string = "Iteration {:0>5}, Cost {:.2f}s".format(
|
||||
self.iteration, now-self.time, end="")
|
||||
for i, (k, v) in enumerate(self.info_dict.items()):
|
||||
if 'scalar' not in k:
|
||||
continue
|
||||
k = k.replace('scalar/', '').replace('/', '_')
|
||||
end = "\n" if i == len(self.info_dict)-1 else ""
|
||||
string += ", {0}={1:.4f}".format(k, np.mean(v), end=end)
|
||||
self.log_info(string)
|
||||
self.reset_time()
|
||||
|
||||
def reset_time(self):
|
||||
self.time = time.time()
|
||||
|
||||
def train_step(self, info, summary):
|
||||
self.iteration += 1
|
||||
self.append(info)
|
||||
if self.iteration % self.log_iter == 0:
|
||||
self.log_training_info()
|
||||
self.flush()
|
||||
self.write_to_tensorboard(summary)
|
||||
|
||||
def log_debug(self, *args, **kwargs):
|
||||
self.logger.debug(*args, **kwargs)
|
||||
|
||||
def log_info(self, *args, **kwargs):
|
||||
self.logger.info(*args, **kwargs)
|
||||
|
||||
def log_warning(self, *args, **kwargs):
|
||||
self.logger.warning(*args, **kwargs)
|
||||
|
||||
|
||||
msg_mgr = MessageManager()
|
||||
noop = NoOp()
|
||||
|
||||
|
||||
def get_msg_mgr():
|
||||
if torch.distributed.get_rank() > 0:
|
||||
return noop
|
||||
else:
|
||||
return msg_mgr
|
||||
Reference in New Issue
Block a user