296 lines
8.4 KiB
Python
296 lines
8.4 KiB
Python
|
import logging
|
||
|
|
import os
|
||
|
|
import pickle
|
||
|
|
import random
|
||
|
|
import shutil
|
||
|
|
import subprocess
|
||
|
|
import SharedArray
|
||
|
|
|
||
|
|
import numpy as np
|
||
|
|
import torch
|
||
|
|
import torch.distributed as dist
|
||
|
|
import torch.multiprocessing as mp
|
||
|
|
|
||
|
|
|
||
|
|
def check_numpy_to_torch(x):
|
||
|
|
if isinstance(x, np.ndarray):
|
||
|
|
return torch.from_numpy(x).float(), True
|
||
|
|
return x, False
|
||
|
|
|
||
|
|
|
||
|
|
def limit_period(val, offset=0.5, period=np.pi):
|
||
|
|
val, is_numpy = check_numpy_to_torch(val)
|
||
|
|
ans = val - torch.floor(val / period + offset) * period
|
||
|
|
return ans.numpy() if is_numpy else ans
|
||
|
|
|
||
|
|
|
||
|
|
def drop_info_with_name(info, name):
|
||
|
|
ret_info = {}
|
||
|
|
keep_indices = [i for i, x in enumerate(info['name']) if x != name]
|
||
|
|
for key in info.keys():
|
||
|
|
ret_info[key] = info[key][keep_indices]
|
||
|
|
return ret_info
|
||
|
|
|
||
|
|
|
||
|
|
def rotate_points_along_z(points, angle):
|
||
|
|
"""
|
||
|
|
Args:
|
||
|
|
points: (B, N, 3 + C)
|
||
|
|
angle: (B), angle along z-axis, angle increases x ==> y
|
||
|
|
Returns:
|
||
|
|
|
||
|
|
"""
|
||
|
|
points, is_numpy = check_numpy_to_torch(points)
|
||
|
|
angle, _ = check_numpy_to_torch(angle)
|
||
|
|
|
||
|
|
cosa = torch.cos(angle)
|
||
|
|
sina = torch.sin(angle)
|
||
|
|
zeros = angle.new_zeros(points.shape[0])
|
||
|
|
ones = angle.new_ones(points.shape[0])
|
||
|
|
rot_matrix = torch.stack((
|
||
|
|
cosa, sina, zeros,
|
||
|
|
-sina, cosa, zeros,
|
||
|
|
zeros, zeros, ones
|
||
|
|
), dim=1).view(-1, 3, 3).float()
|
||
|
|
points_rot = torch.matmul(points[:, :, 0:3], rot_matrix)
|
||
|
|
points_rot = torch.cat((points_rot, points[:, :, 3:]), dim=-1)
|
||
|
|
return points_rot.numpy() if is_numpy else points_rot
|
||
|
|
|
||
|
|
|
||
|
|
def angle2matrix(angle):
|
||
|
|
"""
|
||
|
|
Args:
|
||
|
|
angle: angle along z-axis, angle increases x ==> y
|
||
|
|
Returns:
|
||
|
|
rot_matrix: (3x3 Tensor) rotation matrix
|
||
|
|
"""
|
||
|
|
|
||
|
|
cosa = torch.cos(angle)
|
||
|
|
sina = torch.sin(angle)
|
||
|
|
rot_matrix = torch.tensor([
|
||
|
|
[cosa, -sina, 0],
|
||
|
|
[sina, cosa, 0],
|
||
|
|
[ 0, 0, 1]
|
||
|
|
])
|
||
|
|
return rot_matrix
|
||
|
|
|
||
|
|
|
||
|
|
def mask_points_by_range(points, limit_range):
|
||
|
|
mask = (points[:, 0] >= limit_range[0]) & (points[:, 0] <= limit_range[3]) \
|
||
|
|
& (points[:, 1] >= limit_range[1]) & (points[:, 1] <= limit_range[4])
|
||
|
|
return mask
|
||
|
|
|
||
|
|
|
||
|
|
def get_voxel_centers(voxel_coords, downsample_times, voxel_size, point_cloud_range):
|
||
|
|
"""
|
||
|
|
Args:
|
||
|
|
voxel_coords: (N, 3)
|
||
|
|
downsample_times:
|
||
|
|
voxel_size:
|
||
|
|
point_cloud_range:
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
|
||
|
|
"""
|
||
|
|
assert voxel_coords.shape[1] == 3
|
||
|
|
voxel_centers = voxel_coords[:, [2, 1, 0]].float() # (xyz)
|
||
|
|
voxel_size = torch.tensor(voxel_size, device=voxel_centers.device).float() * downsample_times
|
||
|
|
pc_range = torch.tensor(point_cloud_range[0:3], device=voxel_centers.device).float()
|
||
|
|
voxel_centers = (voxel_centers + 0.5) * voxel_size + pc_range
|
||
|
|
return voxel_centers
|
||
|
|
|
||
|
|
|
||
|
|
def create_logger(log_file=None, rank=0, log_level=logging.INFO):
|
||
|
|
logger = logging.getLogger(__name__)
|
||
|
|
logger.setLevel(log_level if rank == 0 else 'ERROR')
|
||
|
|
formatter = logging.Formatter('%(asctime)s %(levelname)5s %(message)s')
|
||
|
|
console = logging.StreamHandler()
|
||
|
|
console.setLevel(log_level if rank == 0 else 'ERROR')
|
||
|
|
console.setFormatter(formatter)
|
||
|
|
logger.addHandler(console)
|
||
|
|
if log_file is not None:
|
||
|
|
file_handler = logging.FileHandler(filename=log_file)
|
||
|
|
file_handler.setLevel(log_level if rank == 0 else 'ERROR')
|
||
|
|
file_handler.setFormatter(formatter)
|
||
|
|
logger.addHandler(file_handler)
|
||
|
|
logger.propagate = False
|
||
|
|
return logger
|
||
|
|
|
||
|
|
|
||
|
|
def set_random_seed(seed):
|
||
|
|
random.seed(seed)
|
||
|
|
np.random.seed(seed)
|
||
|
|
torch.manual_seed(seed)
|
||
|
|
torch.cuda.manual_seed(seed)
|
||
|
|
torch.backends.cudnn.deterministic = True
|
||
|
|
torch.backends.cudnn.benchmark = False
|
||
|
|
|
||
|
|
|
||
|
|
def worker_init_fn(worker_id, seed=666):
|
||
|
|
if seed is not None:
|
||
|
|
random.seed(seed + worker_id)
|
||
|
|
np.random.seed(seed + worker_id)
|
||
|
|
torch.manual_seed(seed + worker_id)
|
||
|
|
torch.cuda.manual_seed(seed + worker_id)
|
||
|
|
torch.cuda.manual_seed_all(seed + worker_id)
|
||
|
|
|
||
|
|
|
||
|
|
def get_pad_params(desired_size, cur_size):
|
||
|
|
"""
|
||
|
|
Get padding parameters for np.pad function
|
||
|
|
Args:
|
||
|
|
desired_size: int, Desired padded output size
|
||
|
|
cur_size: int, Current size. Should always be less than or equal to cur_size
|
||
|
|
Returns:
|
||
|
|
pad_params: tuple(int), Number of values padded to the edges (before, after)
|
||
|
|
"""
|
||
|
|
assert desired_size >= cur_size
|
||
|
|
|
||
|
|
# Calculate amount to pad
|
||
|
|
diff = desired_size - cur_size
|
||
|
|
pad_params = (0, diff)
|
||
|
|
|
||
|
|
return pad_params
|
||
|
|
|
||
|
|
|
||
|
|
def keep_arrays_by_name(gt_names, used_classes):
|
||
|
|
inds = [i for i, x in enumerate(gt_names) if x in used_classes]
|
||
|
|
inds = np.array(inds, dtype=np.int64)
|
||
|
|
return inds
|
||
|
|
|
||
|
|
|
||
|
|
def init_dist_slurm(tcp_port, local_rank, backend='nccl'):
|
||
|
|
"""
|
||
|
|
modified from https://github.com/open-mmlab/mmdetection
|
||
|
|
Args:
|
||
|
|
tcp_port:
|
||
|
|
backend:
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
|
||
|
|
"""
|
||
|
|
proc_id = int(os.environ['SLURM_PROCID'])
|
||
|
|
ntasks = int(os.environ['SLURM_NTASKS'])
|
||
|
|
node_list = os.environ['SLURM_NODELIST']
|
||
|
|
num_gpus = torch.cuda.device_count()
|
||
|
|
torch.cuda.set_device(proc_id % num_gpus)
|
||
|
|
addr = subprocess.getoutput('scontrol show hostname {} | head -n1'.format(node_list))
|
||
|
|
os.environ['MASTER_PORT'] = str(tcp_port)
|
||
|
|
os.environ['MASTER_ADDR'] = addr
|
||
|
|
os.environ['WORLD_SIZE'] = str(ntasks)
|
||
|
|
os.environ['RANK'] = str(proc_id)
|
||
|
|
dist.init_process_group(backend=backend)
|
||
|
|
|
||
|
|
total_gpus = dist.get_world_size()
|
||
|
|
rank = dist.get_rank()
|
||
|
|
return total_gpus, rank
|
||
|
|
|
||
|
|
|
||
|
|
def init_dist_pytorch(tcp_port, local_rank, backend='nccl'):
|
||
|
|
if mp.get_start_method(allow_none=True) is None:
|
||
|
|
mp.set_start_method('spawn')
|
||
|
|
# os.environ['MASTER_PORT'] = str(tcp_port)
|
||
|
|
# os.environ['MASTER_ADDR'] = 'localhost'
|
||
|
|
num_gpus = torch.cuda.device_count()
|
||
|
|
torch.cuda.set_device(local_rank % num_gpus)
|
||
|
|
|
||
|
|
dist.init_process_group(
|
||
|
|
backend=backend,
|
||
|
|
# init_method='tcp://127.0.0.1:%d' % tcp_port,
|
||
|
|
# rank=local_rank,
|
||
|
|
# world_size=num_gpus
|
||
|
|
)
|
||
|
|
rank = dist.get_rank()
|
||
|
|
return num_gpus, rank
|
||
|
|
|
||
|
|
|
||
|
|
def get_dist_info(return_gpu_per_machine=False):
|
||
|
|
if torch.__version__ < '1.0':
|
||
|
|
initialized = dist._initialized
|
||
|
|
else:
|
||
|
|
if dist.is_available():
|
||
|
|
initialized = dist.is_initialized()
|
||
|
|
else:
|
||
|
|
initialized = False
|
||
|
|
if initialized:
|
||
|
|
rank = dist.get_rank()
|
||
|
|
world_size = dist.get_world_size()
|
||
|
|
else:
|
||
|
|
rank = 0
|
||
|
|
world_size = 1
|
||
|
|
|
||
|
|
if return_gpu_per_machine:
|
||
|
|
gpu_per_machine = torch.cuda.device_count()
|
||
|
|
return rank, world_size, gpu_per_machine
|
||
|
|
|
||
|
|
return rank, world_size
|
||
|
|
|
||
|
|
|
||
|
|
def merge_results_dist(result_part, size, tmpdir):
|
||
|
|
rank, world_size = get_dist_info()
|
||
|
|
os.makedirs(tmpdir, exist_ok=True)
|
||
|
|
|
||
|
|
dist.barrier()
|
||
|
|
pickle.dump(result_part, open(os.path.join(tmpdir, 'result_part_{}.pkl'.format(rank)), 'wb'))
|
||
|
|
dist.barrier()
|
||
|
|
|
||
|
|
if rank != 0:
|
||
|
|
return None
|
||
|
|
|
||
|
|
part_list = []
|
||
|
|
for i in range(world_size):
|
||
|
|
part_file = os.path.join(tmpdir, 'result_part_{}.pkl'.format(i))
|
||
|
|
part_list.append(pickle.load(open(part_file, 'rb')))
|
||
|
|
|
||
|
|
ordered_results = []
|
||
|
|
for res in zip(*part_list):
|
||
|
|
ordered_results.extend(list(res))
|
||
|
|
ordered_results = ordered_results[:size]
|
||
|
|
shutil.rmtree(tmpdir)
|
||
|
|
return ordered_results
|
||
|
|
|
||
|
|
|
||
|
|
def scatter_point_inds(indices, point_inds, shape):
|
||
|
|
ret = -1 * torch.ones(*shape, dtype=point_inds.dtype, device=point_inds.device)
|
||
|
|
ndim = indices.shape[-1]
|
||
|
|
flattened_indices = indices.view(-1, ndim)
|
||
|
|
slices = [flattened_indices[:, i] for i in range(ndim)]
|
||
|
|
ret[slices] = point_inds
|
||
|
|
return ret
|
||
|
|
|
||
|
|
|
||
|
|
def generate_voxel2pinds(sparse_tensor):
|
||
|
|
device = sparse_tensor.indices.device
|
||
|
|
batch_size = sparse_tensor.batch_size
|
||
|
|
spatial_shape = sparse_tensor.spatial_shape
|
||
|
|
indices = sparse_tensor.indices.long()
|
||
|
|
point_indices = torch.arange(indices.shape[0], device=device, dtype=torch.int32)
|
||
|
|
output_shape = [batch_size] + list(spatial_shape)
|
||
|
|
v2pinds_tensor = scatter_point_inds(indices, point_indices, output_shape)
|
||
|
|
return v2pinds_tensor
|
||
|
|
|
||
|
|
|
||
|
|
def sa_create(name, var):
|
||
|
|
x = SharedArray.create(name, var.shape, dtype=var.dtype)
|
||
|
|
x[...] = var[...]
|
||
|
|
x.flags.writeable = False
|
||
|
|
return x
|
||
|
|
|
||
|
|
|
||
|
|
class AverageMeter(object):
|
||
|
|
"""Computes and stores the average and current value"""
|
||
|
|
def __init__(self):
|
||
|
|
self.reset()
|
||
|
|
|
||
|
|
def reset(self):
|
||
|
|
self.val = 0
|
||
|
|
self.avg = 0
|
||
|
|
self.sum = 0
|
||
|
|
self.count = 0
|
||
|
|
|
||
|
|
def update(self, val, n=1):
|
||
|
|
self.val = val
|
||
|
|
self.sum += val * n
|
||
|
|
self.count += n
|
||
|
|
self.avg = self.sum / self.count
|