KPConv-PyTorch/datasets/ModelNet40.py
2020-03-31 15:42:35 -04:00

997 lines
35 KiB
Python

#
#
# 0=================================0
# | Kernel Point Convolutions |
# 0=================================0
#
#
# ----------------------------------------------------------------------------------------------------------------------
#
# Class handling ModelNet40 dataset.
# Implements a Dataset, a Sampler, and a collate_fn
#
# ----------------------------------------------------------------------------------------------------------------------
#
# Hugues THOMAS - 11/06/2018
#
# ----------------------------------------------------------------------------------------------------------------------
#
# Imports and global variables
# \**********************************/
#
# Common libs
import time
import numpy as np
import pickle
import torch
import math
# OS functions
from os import listdir
from os.path import exists, join
# Dataset parent class
from datasets.common import PointCloudDataset
from torch.utils.data import Sampler, get_worker_info
from utils.mayavi_visu import *
from datasets.common import grid_subsampling
from utils.config import bcolors
# ----------------------------------------------------------------------------------------------------------------------
#
# Dataset class definition
# \******************************/
class ModelNet40Dataset(PointCloudDataset):
"""Class to handle Modelnet 40 dataset."""
def __init__(self, config, train=True, orient_correction=True):
"""
This dataset is small enough to be stored in-memory, so load all point clouds here
"""
PointCloudDataset.__init__(self, 'ModelNet40')
############
# Parameters
############
# Dict from labels to names
self.label_to_names = {0: 'airplane',
1: 'bathtub',
2: 'bed',
3: 'bench',
4: 'bookshelf',
5: 'bottle',
6: 'bowl',
7: 'car',
8: 'chair',
9: 'cone',
10: 'cup',
11: 'curtain',
12: 'desk',
13: 'door',
14: 'dresser',
15: 'flower_pot',
16: 'glass_box',
17: 'guitar',
18: 'keyboard',
19: 'lamp',
20: 'laptop',
21: 'mantel',
22: 'monitor',
23: 'night_stand',
24: 'person',
25: 'piano',
26: 'plant',
27: 'radio',
28: 'range_hood',
29: 'sink',
30: 'sofa',
31: 'stairs',
32: 'stool',
33: 'table',
34: 'tent',
35: 'toilet',
36: 'tv_stand',
37: 'vase',
38: 'wardrobe',
39: 'xbox'}
# Initialize a bunch of variables concerning class labels
self.init_labels()
# List of classes ignored during training (can be empty)
self.ignored_labels = np.array([])
# Dataset folder
self.path = '../../Data/ModelNet40'
# Type of task conducted on this dataset
self.dataset_task = 'classification'
# Update number of class and data task in configuration
config.num_classes = self.num_classes
config.dataset_task = self.dataset_task
# Parameters from config
self.config = config
# Training or test set
self.train = train
# Number of models and models used per epoch
if self.train:
self.num_models = 9843
if config.epoch_steps and config.epoch_steps * config.batch_num < self.num_models:
self.epoch_n = config.epoch_steps * config.batch_num
else:
self.epoch_n = self.num_models
else:
self.num_models = 2468
self.epoch_n = min(self.num_models, config.validation_size * config.batch_num)
#############
# Load models
#############
if 0 < self.config.first_subsampling_dl <= 0.01:
raise ValueError('subsampling_parameter too low (should be over 1 cm')
self.input_points, self.input_normals, self.input_labels = self.load_subsampled_clouds(orient_correction)
return
def __len__(self):
"""
Return the length of data here
"""
return self.num_models
def __getitem__(self, idx_list):
"""
The main thread gives a list of indices to load a batch. Each worker is going to work in parallel to load a
different list of indices.
"""
###################
# Gather batch data
###################
tp_list = []
tn_list = []
tl_list = []
ti_list = []
s_list = []
R_list = []
for p_i in idx_list:
# Get points and labels
points = self.input_points[p_i].astype(np.float32)
normals = self.input_normals[p_i].astype(np.float32)
label = self.label_to_idx[self.input_labels[p_i]]
# Data augmentation
points, normals, scale, R = self.augmentation_transform(points, normals)
# Stack batch
tp_list += [points]
tn_list += [normals]
tl_list += [label]
ti_list += [p_i]
s_list += [scale]
R_list += [R]
###################
# Concatenate batch
###################
#show_ModelNet_examples(tp_list, cloud_normals=tn_list)
stacked_points = np.concatenate(tp_list, axis=0)
stacked_normals = np.concatenate(tn_list, axis=0)
labels = np.array(tl_list, dtype=np.int64)
model_inds = np.array(ti_list, dtype=np.int32)
stack_lengths = np.array([tp.shape[0] for tp in tp_list], dtype=np.int32)
scales = np.array(s_list, dtype=np.float32)
rots = np.stack(R_list, axis=0)
# Input features
stacked_features = np.ones_like(stacked_points[:, :1], dtype=np.float32)
if self.config.in_features_dim == 1:
pass
elif self.config.in_features_dim == 4:
stacked_features = np.hstack((stacked_features, stacked_normals))
else:
raise ValueError('Only accepted input dimensions are 1, 4 and 7 (without and with XYZ)')
#######################
# Create network inputs
#######################
#
# Points, neighbors, pooling indices for each layers
#
# Get the whole input list
input_list = self.classification_inputs(stacked_points,
stacked_features,
labels,
stack_lengths)
# Add scale and rotation for testing
input_list += [scales, rots, model_inds]
return input_list
def load_subsampled_clouds(self, orient_correction):
# Restart timer
t0 = time.time()
# Load wanted points if possible
if self.train:
split ='training'
else:
split = 'test'
print('\nLoading {:s} points subsampled at {:.3f}'.format(split, self.config.first_subsampling_dl))
filename = join(self.path, '{:s}_{:.3f}_record.pkl'.format(split, self.config.first_subsampling_dl))
if exists(filename):
with open(filename, 'rb') as file:
input_points, input_normals, input_labels = pickle.load(file)
# Else compute them from original points
else:
# Collect training file names
if self.train:
names = np.loadtxt(join(self.path, 'modelnet40_train.txt'), dtype=np.str)
else:
names = np.loadtxt(join(self.path, 'modelnet40_test.txt'), dtype=np.str)
# Initialize containers
input_points = []
input_normals = []
# Advanced display
N = len(names)
progress_n = 30
fmt_str = '[{:<' + str(progress_n) + '}] {:5.1f}%'
# Collect point clouds
for i, cloud_name in enumerate(names):
# Read points
class_folder = '_'.join(cloud_name.split('_')[:-1])
txt_file = join(self.path, class_folder, cloud_name) + '.txt'
data = np.loadtxt(txt_file, delimiter=',', dtype=np.float32)
# Subsample them
if self.config.first_subsampling_dl > 0:
points, normals = grid_subsampling(data[:, :3],
features=data[:, 3:],
sampleDl=self.config.first_subsampling_dl)
else:
points = data[:, :3]
normals = data[:, 3:]
print('', end='\r')
print(fmt_str.format('#' * ((i * progress_n) // N), 100 * i / N), end='', flush=True)
# Add to list
input_points += [points]
input_normals += [normals]
print('', end='\r')
print(fmt_str.format('#' * progress_n, 100), end='', flush=True)
print()
# Get labels
label_names = ['_'.join(name.split('_')[:-1]) for name in names]
input_labels = np.array([self.name_to_label[name] for name in label_names])
# Save for later use
with open(filename, 'wb') as file:
pickle.dump((input_points,
input_normals,
input_labels), file)
lengths = [p.shape[0] for p in input_points]
sizes = [l * 4 * 6 for l in lengths]
print('{:.1f} MB loaded in {:.1f}s'.format(np.sum(sizes) * 1e-6, time.time() - t0))
if orient_correction:
input_points = [pp[:, [0, 2, 1]] for pp in input_points]
input_normals = [nn[:, [0, 2, 1]] for nn in input_normals]
return input_points, input_normals, input_labels
# ----------------------------------------------------------------------------------------------------------------------
#
# Utility classes definition
# \********************************/
class ModelNet40Sampler(Sampler):
"""Sampler for ModelNet40"""
def __init__(self, dataset: ModelNet40Dataset, use_potential=True, balance_labels=False):
Sampler.__init__(self, dataset)
# Does the sampler use potential for regular sampling
self.use_potential = use_potential
# Should be balance the classes when sampling
self.balance_labels = balance_labels
# Dataset used by the sampler (no copy is made in memory)
self.dataset = dataset
# Create potentials
if self.use_potential:
self.potentials = np.random.rand(len(dataset.input_labels)) * 0.1 + 0.1
else:
self.potentials = None
# Initialize value for batch limit (max number of points per batch).
self.batch_limit = 10000
return
def __iter__(self):
"""
Yield next batch indices here
"""
##########################################
# Initialize the list of generated indices
##########################################
if self.use_potential:
if self.balance_labels:
gen_indices = []
pick_n = self.dataset.epoch_n // self.dataset.num_classes + 1
for i, l in enumerate(self.dataset.label_values):
# Get the potentials of the objects of this class
label_inds = np.where(np.equal(self.dataset.input_labels, l))[0]
class_potentials = self.potentials[label_inds]
# Get the indices to generate thanks to potentials
if pick_n < class_potentials.shape[0]:
pick_indices = np.argpartition(class_potentials, pick_n)[:pick_n]
else:
pick_indices = np.random.permutation(class_potentials.shape[0])
class_indices = label_inds[pick_indices]
gen_indices.append(class_indices)
# Stack the chosen indices of all classes
gen_indices = np.random.permutation(np.hstack(gen_indices))
else:
# Get indices with the minimum potential
if self.dataset.epoch_n < self.potentials.shape[0]:
gen_indices = np.argpartition(self.potentials, self.dataset.epoch_n)[:self.dataset.epoch_n]
else:
gen_indices = np.random.permutation(self.potentials.shape[0])
gen_indices = np.random.permutation(gen_indices)
# Update potentials (Change the order for the next epoch)
self.potentials[gen_indices] = np.ceil(self.potentials[gen_indices])
self.potentials[gen_indices] += np.random.rand(gen_indices.shape[0]) * 0.1 + 0.1
else:
if self.balance_labels:
pick_n = self.dataset.epoch_n // self.dataset.num_classes + 1
gen_indices = []
for l in self.dataset.label_values:
label_inds = np.where(np.equal(self.dataset.input_labels, l))[0]
rand_inds = np.random.choice(label_inds, size=pick_n, replace=True)
gen_indices += [rand_inds]
gen_indices = np.random.permutation(np.hstack(gen_indices))
else:
gen_indices = np.random.permutation(self.dataset.num_models)[:self.dataset.epoch_n]
################
# Generator loop
################
# Initialize concatenation lists
ti_list = []
batch_n = 0
# Generator loop
for p_i in gen_indices:
# Size of picked cloud
n = self.dataset.input_points[p_i].shape[0]
# In case batch is full, yield it and reset it
if batch_n + n > self.batch_limit and batch_n > 0:
yield np.array(ti_list, dtype=np.int32)
ti_list = []
batch_n = 0
# Add data to current batch
ti_list += [p_i]
# Update batch size
batch_n += n
yield np.array(ti_list, dtype=np.int32)
return 0
def __len__(self):
"""
The number of yielded samples is variable
"""
return None
def calibration(self, dataloader, untouched_ratio=0.9, verbose=False):
"""
Method performing batch and neighbors calibration.
Batch calibration: Set "batch_limit" (the maximum number of points allowed in every batch) so that the
average batch size (number of stacked pointclouds) is the one asked.
Neighbors calibration: Set the "neighborhood_limits" (the maximum number of neighbors allowed in convolutions)
so that 90% of the neighborhoods remain untouched. There is a limit for each layer.
"""
##############################
# Previously saved calibration
##############################
print('\nStarting Calibration (use verbose=True for more details)')
t0 = time.time()
redo = False
# Batch limit
# ***********
# Load batch_limit dictionary
batch_lim_file = join(self.dataset.path, 'batch_limits.pkl')
if exists(batch_lim_file):
with open(batch_lim_file, 'rb') as file:
batch_lim_dict = pickle.load(file)
else:
batch_lim_dict = {}
# Check if the batch limit associated with current parameters exists
key = '{:.3f}_{:d}'.format(self.dataset.config.first_subsampling_dl,
self.dataset.config.batch_num)
if key in batch_lim_dict:
self.batch_limit = batch_lim_dict[key]
else:
redo = True
if verbose:
print('\nPrevious calibration found:')
print('Check batch limit dictionary')
if key in batch_lim_dict:
color = bcolors.OKGREEN
v = str(int(batch_lim_dict[key]))
else:
color = bcolors.FAIL
v = '?'
print('{:}\"{:s}\": {:s}{:}'.format(color, key, v, bcolors.ENDC))
# Neighbors limit
# ***************
# Load neighb_limits dictionary
neighb_lim_file = join(self.dataset.path, 'neighbors_limits.pkl')
if exists(neighb_lim_file):
with open(neighb_lim_file, 'rb') as file:
neighb_lim_dict = pickle.load(file)
else:
neighb_lim_dict = {}
# Check if the limit associated with current parameters exists (for each layer)
neighb_limits = []
for layer_ind in range(self.dataset.config.num_layers):
dl = self.dataset.config.first_subsampling_dl * (2**layer_ind)
if self.dataset.config.deform_layers[layer_ind]:
r = dl * self.dataset.config.deform_radius
else:
r = dl * self.dataset.config.conv_radius
key = '{:.3f}_{:.3f}'.format(dl, r)
if key in neighb_lim_dict:
neighb_limits += [neighb_lim_dict[key]]
if len(neighb_limits) == self.dataset.config.num_layers:
self.dataset.neighborhood_limits = neighb_limits
else:
redo = True
if verbose:
print('Check neighbors limit dictionary')
for layer_ind in range(self.dataset.config.num_layers):
dl = self.dataset.config.first_subsampling_dl * (2**layer_ind)
if self.dataset.config.deform_layers[layer_ind]:
r = dl * self.dataset.config.deform_radius
else:
r = dl * self.dataset.config.conv_radius
key = '{:.3f}_{:.3f}'.format(dl, r)
if key in neighb_lim_dict:
color = bcolors.OKGREEN
v = str(neighb_lim_dict[key])
else:
color = bcolors.FAIL
v = '?'
print('{:}\"{:s}\": {:s}{:}'.format(color, key, v, bcolors.ENDC))
if redo:
############################
# Neighbors calib parameters
############################
# From config parameter, compute higher bound of neighbors number in a neighborhood
hist_n = int(np.ceil(4 / 3 * np.pi * (self.dataset.config.conv_radius + 1) ** 3))
# Histogram of neighborhood sizes
neighb_hists = np.zeros((self.dataset.config.num_layers, hist_n), dtype=np.int32)
########################
# Batch calib parameters
########################
# Estimated average batch size and target value
estim_b = 0
target_b = self.dataset.config.batch_num
# Calibration parameters
low_pass_T = 10
Kp = 100.0
finer = False
# Convergence parameters
smooth_errors = []
converge_threshold = 0.1
# Loop parameters
last_display = time.time()
i = 0
breaking = False
#####################
# Perform calibration
#####################
for epoch in range(10):
for batch_i, batch in enumerate(dataloader):
# Update neighborhood histogram
counts = [np.sum(neighb_mat.numpy() < neighb_mat.shape[0], axis=1) for neighb_mat in batch.neighbors]
hists = [np.bincount(c, minlength=hist_n)[:hist_n] for c in counts]
neighb_hists += np.vstack(hists)
# batch length
b = len(batch.labels)
# Update estim_b (low pass filter)
estim_b += (b - estim_b) / low_pass_T
# Estimate error (noisy)
error = target_b - b
# Save smooth errors for convergene check
smooth_errors.append(target_b - estim_b)
if len(smooth_errors) > 10:
smooth_errors = smooth_errors[1:]
# Update batch limit with P controller
self.batch_limit += Kp * error
# finer low pass filter when closing in
if not finer and np.abs(estim_b - target_b) < 1:
low_pass_T = 100
finer = True
# Convergence
if finer and np.max(np.abs(smooth_errors)) < converge_threshold:
breaking = True
break
i += 1
t = time.time()
# Console display (only one per second)
if verbose and (t - last_display) > 1.0:
last_display = t
message = 'Step {:5d} estim_b ={:5.2f} batch_limit ={:7d}'
print(message.format(i,
estim_b,
int(self.batch_limit)))
if breaking:
break
# Use collected neighbor histogram to get neighbors limit
cumsum = np.cumsum(neighb_hists.T, axis=0)
percentiles = np.sum(cumsum < (untouched_ratio * cumsum[hist_n - 1, :]), axis=0)
self.dataset.neighborhood_limits = percentiles
if verbose:
# Crop histogram
while np.sum(neighb_hists[:, -1]) == 0:
neighb_hists = neighb_hists[:, :-1]
hist_n = neighb_hists.shape[1]
print('\n**************************************************\n')
line0 = 'neighbors_num '
for layer in range(neighb_hists.shape[0]):
line0 += '| layer {:2d} '.format(layer)
print(line0)
for neighb_size in range(hist_n):
line0 = ' {:4d} '.format(neighb_size)
for layer in range(neighb_hists.shape[0]):
if neighb_size > percentiles[layer]:
color = bcolors.FAIL
else:
color = bcolors.OKGREEN
line0 += '|{:}{:10d}{:} '.format(color,
neighb_hists[layer, neighb_size],
bcolors.ENDC)
print(line0)
print('\n**************************************************\n')
print('\nchosen neighbors limits: ', percentiles)
print()
# Save batch_limit dictionary
key = '{:.3f}_{:d}'.format(self.dataset.config.first_subsampling_dl,
self.dataset.config.batch_num)
batch_lim_dict[key] = self.batch_limit
with open(batch_lim_file, 'wb') as file:
pickle.dump(batch_lim_dict, file)
# Save neighb_limit dictionary
for layer_ind in range(self.dataset.config.num_layers):
dl = self.dataset.config.first_subsampling_dl * (2 ** layer_ind)
if self.dataset.config.deform_layers[layer_ind]:
r = dl * self.dataset.config.deform_radius
else:
r = dl * self.dataset.config.conv_radius
key = '{:.3f}_{:.3f}'.format(dl, r)
neighb_lim_dict[key] = self.dataset.neighborhood_limits[layer_ind]
with open(neighb_lim_file, 'wb') as file:
pickle.dump(neighb_lim_dict, file)
print('Calibration done in {:.1f}s\n'.format(time.time() - t0))
return
class ModelNet40CustomBatch:
"""Custom batch definition with memory pinning for ModelNet40"""
def __init__(self, input_list):
# Get rid of batch dimension
input_list = input_list[0]
# Number of layers
L = (len(input_list) - 5) // 4
# Extract input tensors from the list of numpy array
ind = 0
self.points = [torch.from_numpy(nparray) for nparray in input_list[ind:ind+L]]
ind += L
self.neighbors = [torch.from_numpy(nparray) for nparray in input_list[ind:ind+L]]
ind += L
self.pools = [torch.from_numpy(nparray) for nparray in input_list[ind:ind+L]]
ind += L
self.lengths = [torch.from_numpy(nparray) for nparray in input_list[ind:ind+L]]
ind += L
self.features = torch.from_numpy(input_list[ind])
ind += 1
self.labels = torch.from_numpy(input_list[ind])
ind += 1
self.scales = torch.from_numpy(input_list[ind])
ind += 1
self.rots = torch.from_numpy(input_list[ind])
ind += 1
self.model_inds = torch.from_numpy(input_list[ind])
return
def pin_memory(self):
"""
Manual pinning of the memory
"""
self.points = [in_tensor.pin_memory() for in_tensor in self.points]
self.neighbors = [in_tensor.pin_memory() for in_tensor in self.neighbors]
self.pools = [in_tensor.pin_memory() for in_tensor in self.pools]
self.lengths = [in_tensor.pin_memory() for in_tensor in self.lengths]
self.features = self.features.pin_memory()
self.labels = self.labels.pin_memory()
self.scales = self.scales.pin_memory()
self.rots = self.rots.pin_memory()
self.model_inds = self.model_inds.pin_memory()
return self
def to(self, device):
self.points = [in_tensor.to(device) for in_tensor in self.points]
self.neighbors = [in_tensor.to(device) for in_tensor in self.neighbors]
self.pools = [in_tensor.to(device) for in_tensor in self.pools]
self.lengths = [in_tensor.to(device) for in_tensor in self.lengths]
self.features = self.features.to(device)
self.labels = self.labels.to(device)
self.scales = self.scales.to(device)
self.rots = self.rots.to(device)
self.model_inds = self.model_inds.to(device)
return self
def unstack_points(self, layer=None):
"""Unstack the points"""
return self.unstack_elements('points', layer)
def unstack_neighbors(self, layer=None):
"""Unstack the neighbors indices"""
return self.unstack_elements('neighbors', layer)
def unstack_pools(self, layer=None):
"""Unstack the pooling indices"""
return self.unstack_elements('pools', layer)
def unstack_elements(self, element_name, layer=None, to_numpy=True):
"""
Return a list of the stacked elements in the batch at a certain layer. If no layer is given, then return all
layers
"""
if element_name == 'points':
elements = self.points
elif element_name == 'neighbors':
elements = self.neighbors
elif element_name == 'pools':
elements = self.pools[:-1]
else:
raise ValueError('Unknown element name: {:s}'.format(element_name))
all_p_list = []
for layer_i, layer_elems in enumerate(elements):
if layer is None or layer == layer_i:
i0 = 0
p_list = []
if element_name == 'pools':
lengths = self.lengths[layer_i+1]
else:
lengths = self.lengths[layer_i]
for b_i, length in enumerate(lengths):
elem = layer_elems[i0:i0 + length]
if element_name == 'neighbors':
elem[elem >= self.points[layer_i].shape[0]] = -1
elem[elem >= 0] -= i0
elif element_name == 'pools':
elem[elem >= self.points[layer_i].shape[0]] = -1
elem[elem >= 0] -= torch.sum(self.lengths[layer_i][:b_i])
i0 += length
if to_numpy:
p_list.append(elem.numpy())
else:
p_list.append(elem)
if layer == layer_i:
return p_list
all_p_list.append(p_list)
return all_p_list
def ModelNet40Collate(batch_data):
return ModelNet40CustomBatch(batch_data)
class ModelNet40WorkerInitDebug():
"""Callable class that Initializes workers."""
def __init__(self, dataset):
self.dataset = dataset
return
def __call__(self, worker_id):
# Print workers info
worker_info = get_worker_info()
print(worker_info)
# Get associated dataset
dataset = worker_info.dataset # the dataset copy in this worker process
# In windows, each worker has its own copy of the dataset. In Linux, this is shared in memory
print(dataset.input_labels.__array_interface__['data'])
print(worker_info.dataset.input_labels.__array_interface__['data'])
print(self.dataset.input_labels.__array_interface__['data'])
# configure the dataset to only process the split workload
return
# ----------------------------------------------------------------------------------------------------------------------
#
# Debug functions
# \*********************/
def debug_sampling(dataset, sampler, loader):
"""Shows which labels are sampled according to strategy chosen"""
label_sum = np.zeros((dataset.num_classes), dtype=np.int32)
for epoch in range(10):
for batch_i, (points, normals, labels, indices, in_sizes) in enumerate(loader):
# print(batch_i, tuple(points.shape), tuple(normals.shape), labels, indices, in_sizes)
label_sum += np.bincount(labels.numpy(), minlength=dataset.num_classes)
print(label_sum)
#print(sampler.potentials[:6])
print('******************')
print('*******************************************')
_, counts = np.unique(dataset.input_labels, return_counts=True)
print(counts)
def debug_timing(dataset, sampler, loader):
"""Timing of generator function"""
t = [time.time()]
last_display = time.time()
mean_dt = np.zeros(2)
estim_b = dataset.config.batch_num
for epoch in range(10):
for batch_i, batch in enumerate(loader):
# print(batch_i, tuple(points.shape), tuple(normals.shape), labels, indices, in_sizes)
# New time
t = t[-1:]
t += [time.time()]
# Update estim_b (low pass filter)
estim_b += (len(batch.labels) - estim_b) / 100
# Pause simulating computations
time.sleep(0.050)
t += [time.time()]
# Average timing
mean_dt = 0.9 * mean_dt + 0.1 * (np.array(t[1:]) - np.array(t[:-1]))
# Console display (only one per second)
if (t[-1] - last_display) > -1.0:
last_display = t[-1]
message = 'Step {:08d} -> (ms/batch) {:8.2f} {:8.2f} / batch = {:.2f}'
print(message.format(batch_i,
1000 * mean_dt[0],
1000 * mean_dt[1],
estim_b))
print('************* Epoch ended *************')
_, counts = np.unique(dataset.input_labels, return_counts=True)
print(counts)
def debug_show_clouds(dataset, sampler, loader):
for epoch in range(10):
clouds = []
cloud_normals = []
cloud_labels = []
L = dataset.config.num_layers
for batch_i, batch in enumerate(loader):
# Print characteristics of input tensors
print('\nPoints tensors')
for i in range(L):
print(batch.points[i].dtype, batch.points[i].shape)
print('\nNeigbors tensors')
for i in range(L):
print(batch.neighbors[i].dtype, batch.neighbors[i].shape)
print('\nPools tensors')
for i in range(L):
print(batch.pools[i].dtype, batch.pools[i].shape)
print('\nStack lengths')
for i in range(L):
print(batch.lengths[i].dtype, batch.lengths[i].shape)
print('\nFeatures')
print(batch.features.dtype, batch.features.shape)
print('\nLabels')
print(batch.labels.dtype, batch.labels.shape)
print('\nAugment Scales')
print(batch.scales.dtype, batch.scales.shape)
print('\nAugment Rotations')
print(batch.rots.dtype, batch.rots.shape)
print('\nModel indices')
print(batch.model_inds.dtype, batch.model_inds.shape)
print('\nAre input tensors pinned')
print(batch.neighbors[0].is_pinned())
print(batch.neighbors[-1].is_pinned())
print(batch.points[0].is_pinned())
print(batch.points[-1].is_pinned())
print(batch.labels.is_pinned())
print(batch.scales.is_pinned())
print(batch.rots.is_pinned())
print(batch.model_inds.is_pinned())
show_input_batch(batch)
print('*******************************************')
_, counts = np.unique(dataset.input_labels, return_counts=True)
print(counts)
def debug_batch_and_neighbors_calib(dataset, sampler, loader):
"""Timing of generator function"""
t = [time.time()]
last_display = time.time()
mean_dt = np.zeros(2)
for epoch in range(10):
for batch_i, input_list in enumerate(loader):
# print(batch_i, tuple(points.shape), tuple(normals.shape), labels, indices, in_sizes)
# New time
t = t[-1:]
t += [time.time()]
# Pause simulating computations
time.sleep(0.01)
t += [time.time()]
# Average timing
mean_dt = 0.9 * mean_dt + 0.1 * (np.array(t[1:]) - np.array(t[:-1]))
# Console display (only one per second)
if (t[-1] - last_display) > 1.0:
last_display = t[-1]
message = 'Step {:08d} -> Average timings (ms/batch) {:8.2f} {:8.2f} '
print(message.format(batch_i,
1000 * mean_dt[0],
1000 * mean_dt[1]))
print('************* Epoch ended *************')
_, counts = np.unique(dataset.input_labels, return_counts=True)
print(counts)