def forward(self):
blob = {}
blob['data'] = np.reshape(self.file.root.data[self.index], (1,) + (self.N,) * self.dim + (1,))
if not self.is_testing:
blob['labels'] = np.reshape(self.file.root.label[self.index], (1,) + (self.N,) * self.dim)
blob['voxels'], blob['voxels_value'] = extract_voxels(blob['data'][0, ..., 0])
blob['entries'] = [self.index]
self.index = (self.index + 1) % self.n
return blob
def cluster(cfg, blob, results, index, name='cluster', directory=None):
"""
Ad-hoc clustering algorithm. Can use UResNet predictions as a mask for
to cluster track and shower separately, if results includes `predictions`
key. Erases a 7x7 window around each point predicted by PPN in the data,
then applies DBSCAN algorithm to perform rough clustering of track/shower
instances.
"""
data = blob['data']
WINDOW_SIZE = 7
# Hide window around each proposal
for p in results['im_proposals']:
p1 = (p - WINDOW_SIZE / 2).astype(int)
p2 = (p + WINDOW_SIZE / 2).astype(int)
if cfg.DATA_3D:
data[0, p1[0]:p2[0], p1[1]:p2[1], p1[2]:p2[2], 0] = 0.0
else:
data[0, p1[0]:p2[0], p1[1]:p2[1], 0] = 0.0
if 'predictions' in results: # UResNet mask
predictions = results['predictions'][0, ...]
predictions[data[0, ..., 0] == 0.0] = 0.0 # mask with data
track_voxels = np.argwhere(predictions == 1) # track
shower_voxels = np.argwhere(predictions == 2) # track
db_track, db_shower = [], []
if track_voxels.shape[0]:
db_track = DBSCAN(eps=2.5 if cfg.DATA_3D else 2.0,
min_samples=10).fit_predict(track_voxels)
if shower_voxels.shape[0]:
db_shower = DBSCAN(eps=2.5 if cfg.DATA_3D else 2.0,
min_samples=10).fit_predict(shower_voxels)
db_shower = db_shower + len(np.unique(db_track)) # offset labels
voxels = np.concatenate([track_voxels, shower_voxels], axis=0)
voxels = np.flip(voxels, axis=1)
db = np.concatenate([db_track, db_shower], axis=0)
else:
voxels, _ = extract_voxels(data[0, ..., 0])
voxels = np.flip(voxels, axis=1)
db = DBSCAN(eps=2.5 if cfg.DATA_3D else 2.0,
min_samples=10).fit_predict(voxels)
print("Clusters: ", np.unique(db))
new_blob = {}
new_blob['data'] = data
new_blob['voxels'] = voxels
new_blob['voxels_value'] = db
# display_uresnet(new_blob, cfg,
# compute_voxels=False,
# directory=directory,
# vmin=np.amin(db),
# vmax=np.amax(db),
# index=index)
kwargs = {}
if cfg.DATA_3D:
kwargs['projection'] = '3d'
display_blob(new_blob,
cfg,
directory=directory,
index=index,
cmap='tab10',
**kwargs)
def test(cfg, net):
_, data = get_data(cfg)
# Initialize the network the right way
# net.train and net.eval account for differences in dropout/batch norm
# during training and testing
net.eval().cuda()
metrics = {'acc_all': [], 'acc_nonzero': [], 'loss': []}
metrics_mean = {'acc_all': [], 'acc_nonzero': [], 'loss': []}
metrics_std = {'acc_all': [], 'acc_nonzero': [], 'loss': []}
durations_mean = {'cuda': [], 'loss': [], 'forward': [], 'acc': []}
durations_std = {'cuda': [], 'loss': [], 'forward': [], 'acc': []}
# Only enable gradients if we are training
# with torch.set_grad_enabled(is_training):
durations, durations_cuda, durations_loss, durations_acc = [], [], [], []
steps = []
print('Listing weights...')
weights = glob.glob(os.path.join(cfg.WEIGHTS_FILE_BASE, "*.ckpt"))
weights.sort()
print('Done.')
blobs = []
print('Fetch data...')
for i in range(cfg.MAX_STEPS):
print("%d/%d" % (i, cfg.MAX_STEPS))
blob = data.forward()
blob.pop('data')
blob['label_voxels'], blob['label_values'] = extract_voxels(
blob['labels'])
blob.pop('labels')
blobs.append(blob)
print('Done.')
for w in weights:
step = int(re.findall(r'model-(\d+)', w)[0])
steps.append(step)
print('Restoring weights from %s...' % w)
with open(w, 'rb') as f:
checkpoint = torch.load(f)
net.load_state_dict(checkpoint['state_dict'])
print('Done.')
for i, blob in enumerate(blobs): # FIXME
print("Step %d/%d" % (i, cfg.MAX_STEPS))
if sparse:
start = time.time()
coords = torch.from_numpy(blob['voxels']).cuda()
features = torch.from_numpy(
np.reshape(blob['voxels_value'], (-1, 1))).cuda()
label_voxels, labels = blob['label_voxels'], blob[
'label_values']
labels = torch.from_numpy(labels).cuda().type(
torch.cuda.LongTensor)
end = time.time()
durations_cuda.append(end - start)
start = time.time()
predictions_raw = net(coords,
features) # size N_voxels x num_classes
end = time.time()
durations.append(end - start)
else:
start = time.time()
image = torch.from_numpy(np.moveaxis(blob['data'], -1,
1)).cuda()
labels = torch.from_numpy(blob['labels']).cuda().type(
torch.cuda.LongTensor)
end = time.time()
durations_cuda.append(end - start)
start = time.time()
predictions_raw = net(image)
end = time.time()
durations.append(end - start)
start = time.time()
loss = criterion(predictions_raw, labels)
end = time.time()
durations_loss.append(end - start)
metrics['loss'].append(loss.item())
print("\tLoss = ", metrics['loss'][-1])
# Accuracy
start = time.time()
predicted_labels = torch.argmax(predictions_raw, dim=1)
acc_all = (predicted_labels == labels).sum().item() / float(
labels.numel())
nonzero_px = labels > 0
nonzero_prediction = predicted_labels[nonzero_px]
nonzero_label = labels[nonzero_px]
acc_nonzero = (nonzero_prediction
== nonzero_label).sum().item() / float(
nonzero_label.numel())
end = time.time()
durations_acc.append(end - start)
metrics['acc_all'].append(acc_all)
metrics['acc_nonzero'].append(acc_nonzero)
print("\tAccuracy = ", metrics['acc_all'][-1],
" - Nonzero accuracy = ", metrics['acc_nonzero'][-1])
metrics_mean['loss'].append(np.array(metrics['loss']).mean())
metrics_std['loss'].append(np.array(metrics['loss']).std())
metrics_mean['acc_all'].append(np.array(metrics['acc_all']).mean())
metrics_std['acc_all'].append(np.array(metrics['acc_all']).std())
metrics_mean['acc_nonzero'].append(
np.array(metrics['acc_nonzero']).mean())
metrics_std['acc_nonzero'].append(
np.array(metrics['acc_nonzero']).std())
durations_mean['cuda'].append(np.array(durations_cuda).mean())
durations_std['cuda'].append(np.array(durations_cuda).std())
durations_mean['loss'].append(np.array(durations_loss).mean())
durations_std['loss'].append(np.array(durations_loss).std())
durations_mean['forward'].append(np.array(durations).mean())
durations_std['forward'].append(np.array(durations).std())
durations_mean['acc'].append(np.array(durations_acc).mean())
durations_std['acc'].append(np.array(durations_acc).std())
durations, durations_cuda, durations_loss, durations_acc = [], [], [], []
metrics = {'acc_all': [], 'acc_nonzero': [], 'loss': []}
print('Mean cuda duration = %f s' % durations_mean['cuda'][-1])
print('Mean loss duration = %f s' % durations_mean['loss'][-1])
print('Mean acc duration = %f s' % durations_mean['acc'][-1])
print('Mean forward duration = %f s' % durations_mean['forward'][-1])
print('Mean acc = %f s' % metrics_mean['acc_nonzero'][-1])
np.savetxt(os.path.join(cfg.OUTPUT_DIR, 'steps_%d.csv' % step),
steps,
delimiter=',')
for attr in metrics:
np.savetxt(os.path.join(cfg.OUTPUT_DIR,
'%s_mean_%d.csv' % (attr, step)),
metrics_mean[attr],
delimiter=',')
np.savetxt(os.path.join(cfg.OUTPUT_DIR,
'%s_std_%d.csv' % (attr, step)),
metrics_std[attr],
delimiter=',')
for attr in durations_mean:
np.savetxt(os.path.join(cfg.OUTPUT_DIR,
'durations_%s_mean_%d.csv' % (attr, step)),
durations_mean[attr],
delimiter=',')
np.savetxt(os.path.join(cfg.OUTPUT_DIR,
'durations_%s_std_%d.csv' % (attr, step)),
durations_std[attr],
delimiter=',')
def train_demo(cfg, net, criterion, optimizer, lr_scheduler):
# Data generator
train_data, test_data = get_data(cfg)
if is_training:
data = train_data
else:
data = test_data
# Initialize the network the right way
# net.train and net.eval account for differences in dropout/batch norm
# during training and testing
start = 0
if is_training:
net.train().cuda()
else:
net.eval().cuda()
if cfg.WEIGHTS_FILE_BASE is not None and cfg.WEIGHTS_FILE_BASE != '':
print('Restoring weights from %s...' % cfg.WEIGHTS_FILE_BASE)
with open(cfg.WEIGHTS_FILE_BASE, 'rb') as f:
checkpoint = torch.load(f)
net.load_state_dict(checkpoint['state_dict'])
# print(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
start = checkpoint['epoch'] + 1
print('Done.')
print('Done.')
metrics = {'acc_all': [], 'acc_nonzero': [], 'loss': []}
# Only enable gradients if we are training
# with torch.set_grad_enabled(is_training):
durations = []
for i in range(cfg.MAX_STEPS): # use with torch.no_grad() for test network
# Check parameters for nan
# print('Check for nan...')
# had_nan = False
# for p in net.parameters():
# if torch.isnan(p).any():
# print(i, p)
# had_nan = True
#
# for name in net.state_dict():
# tensor = net.state_dict()[name]
# if name == 'sparseModel.2.4.1.2.4.1.2.4.1.2.4.1.2.0.1.0.runningVar':
# print(i, name, tensor)
# if torch.isnan(tensor).any():
# print(i, name, tensor)
# had_nan = True
# if had_nan:
# break
# print('Done.')
# inputs, label = dataloader(i)
print("Step %d/%d" % (i, cfg.MAX_STEPS))
blob = data.forward()
print(blob['voxels'].shape, blob['voxels_value'].shape,
blob['data'].shape, blob['labels'].shape)
if sparse:
coords = torch.from_numpy(blob['voxels']).cuda()
features = torch.from_numpy(
np.reshape(blob['voxels_value'], (-1, 1))).cuda()
# print(coords.type(), features.type())
start = time.time()
predictions_raw = net(coords,
features) # size N_voxels x num_classes
end = time.time()
durations.append(end - start)
# print(predictions_raw.size())
label_voxels, labels = extract_voxels(blob['labels'])
labels = torch.from_numpy(labels).cuda().type(
torch.cuda.LongTensor)
# print(labels, label_voxels, blob['voxels'])
# print(net.parameters())
else:
image = torch.from_numpy(np.moveaxis(blob['data'], -1, 1)).cuda()
labels = torch.from_numpy(blob['labels']).cuda().type(
torch.cuda.LongTensor)
start = time.time()
predictions_raw = net(image)
end = time.time()
durations.append(end - start)
loss = criterion(predictions_raw, labels)
if is_training:
lr_scheduler.step() # Decay learning rate
optimizer.zero_grad() # Clear previous gradients
loss.backward() # Compute gradients of all variables wrt loss
nn.utils.clip_grad_norm_(net.parameters(), 1.0) # Clip gradient
optimizer.step() # update using computed gradients
metrics['loss'].append(loss.item())
print("\tLoss = ", metrics['loss'][-1])
# Accuracy
predicted_labels = torch.argmax(predictions_raw, dim=1)
acc_all = (predicted_labels == labels).sum().item() / float(
labels.numel())
nonzero_px = labels > 0
nonzero_prediction = predicted_labels[nonzero_px]
nonzero_label = labels[nonzero_px]
acc_nonzero = (nonzero_prediction
== nonzero_label).sum().item() / float(
nonzero_label.numel())
metrics['acc_all'].append(acc_all)
metrics['acc_nonzero'].append(acc_nonzero)
print("\tAccuracy = ", metrics['acc_all'][-1],
" - Nonzero accuracy = ", metrics['acc_nonzero'][-1])
if is_training and i % 100 == 0:
for attr in metrics:
np.savetxt(os.path.join(cfg.OUTPUT_DIR,
'%s_%d.csv' % (attr, i)),
metrics[attr],
delimiter=',')
# metrics[attr] = []
if is_training and i % 100 == 0:
filename = os.path.join(cfg.OUTPUT_DIR, 'model-%d.ckpt' % i)
# with open(filename, 'wb'):
torch.save(
{
'epoch': i,
'state_dict': net.state_dict(),
'optimizer': optimizer.state_dict(),
'lr_scheduler': lr_scheduler.state_dict()
}, filename)
if not is_training:
print('Display...')
if sparse:
final_predictions = np.zeros(
(1, cfg.IMAGE_SIZE, cfg.IMAGE_SIZE, cfg.IMAGE_SIZE))
indices = label_voxels.T
final_predictions[
0, indices[0], indices[1],
indices[2]] = predicted_labels.cpu().data.numpy()
display_uresnet(blob,
cfg,
index=i,
predictions=final_predictions)
else:
display_uresnet(
blob,
cfg,
index=i,
predictions=predicted_labels.cpu().data.numpy())
print('Done.')
print("Average duration = %f s" % np.array(durations).mean())
def reconcile(self, batch_results, patch_centers, patch_sizes):
"""
Reconcile slices result together
using batch_results, batch_blobs, patch_centers and patch_sizes
"""
final_results = {}
if len(batch_results) == 0: # Empty batch
return final_results
# UResNet predictions
if 'predictions' and 'scores' and 'softmax' in batch_results[0]:
final_voxels = np.array([], dtype=np.int32).reshape(
0, 3) # Shape N_voxels x dim
final_scores = np.array([], dtype=np.float32).reshape(
0, self.cfg.NUM_CLASSES) # Shape N_voxels x num_classes
final_counts = np.array([], dtype=np.int32).reshape(
0, ) # Shape N_voxels x 1
for i, result in enumerate(batch_results):
# Extract voxel and voxel values
# Shape N_voxels x dim
v, values = extract_voxels(result['predictions'])
# Extract corresponding softmax scores
# Shape N_voxels x num_classes
scores = result['softmax'][v[:, 0], v[:, 1], v[:, 2], :]
# Restore original blob coordinates
v = (v + np.flipud(patch_centers[i]) -
patch_sizes[i] / 2.0).astype(np.int64)
v = np.clip(v, 0, self.cfg.IMAGE_SIZE - 1)
# indices are indices of the *first* occurrences of the unique values
# hence for doublons they are indices in final_voxels
# We assume the only overlap that can occur is between
# final_voxels and v, not inside these arrays themselves
n = final_voxels.shape[0]
final_voxels, indices, counts = np.unique(np.concatenate(
[final_voxels, v], axis=0),
axis=0,
return_index=True,
return_counts=True)
final_scores = np.concatenate([final_scores, scores],
axis=0)[indices]
lower_indices = indices[indices < n]
upper_indices = indices[indices >= n]
final_counts[lower_indices] += counts[lower_indices] - 1
final_counts = np.concatenate(
[final_counts,
np.ones((upper_indices.shape[0], ))],
axis=0)
final_scores = final_scores / final_counts[:, np.
newaxis] # Compute average
final_predictions = np.argmax(final_scores, axis=1)
final_results['predictions'] = np.zeros(
(self.cfg.IMAGE_SIZE, ) * 3)
final_results['predictions'][final_voxels.T[0], final_voxels.T[1],
final_voxels.T[2]] = final_predictions
final_results['scores'] = np.zeros((self.cfg.IMAGE_SIZE, ) * 3)
final_results['scores'][final_voxels.T[0], final_voxels.T[1],
final_voxels.T[2]] = final_scores[
np.arange(final_scores.shape[0]),
final_predictions]
final_results['softmax'] = np.zeros((self.cfg.IMAGE_SIZE, ) * 3 +
(self.cfg.NUM_CLASSES, ))
final_results['softmax'][final_voxels.T[0], final_voxels.T[1],
final_voxels.T[2], :] = final_scores
final_results['predictions'] = final_results['predictions'][
np.newaxis, ...]
# PPN
if 'im_proposals' and 'im_scores' and 'im_labels' and 'rois' in batch_results[
0]:
# print(batch_results[0]['im_proposals'].shape, batch_results[0]['im_scores'].shape, batch_results[0]['im_labels'].shape, batch_results[0]['rois'].shape)
final_im_proposals = np.array([], dtype=np.float32).reshape(0, 3)
final_im_scores = np.array([], dtype=np.float32).reshape(0, )
final_im_labels = np.array([], dtype=np.int32).reshape(0, )
final_rois = np.array([], dtype=np.float32).reshape(0, 3)
for i, result in enumerate(batch_results):
im_proposals = result['im_proposals'] + np.flipud(
patch_centers[i]) - patch_sizes[i] / 2.0
im_proposals = np.clip(im_proposals, 0,
self.cfg.IMAGE_SIZE - 1)
# print(final_im_proposals, im_proposals)
final_im_proposals = np.concatenate(
[final_im_proposals, im_proposals], axis=0)
final_im_scores = np.concatenate(
[final_im_scores, result['im_scores']], axis=0)
final_im_labels = np.concatenate(
[final_im_labels, result['im_labels']], axis=0)
rois = result['rois'] + (np.flipud(patch_centers[i]) -
patch_sizes[i] / 2.0) / (
self.cfg.dim1 * self.cfg.dim2)
rois = np.clip(rois, 0, self.cfg.IMAGE_SIZE - 1)
final_rois = np.concatenate([final_rois, rois], axis=0)
final_results['im_proposals'] = np.array(final_im_proposals)
final_results['im_scores'] = np.array(final_im_scores)
final_results['im_labels'] = np.array(final_im_labels)
final_results['rois'] = np.array(final_rois)
# Try thresholding
# index = np.where(final_results['im_scores'] > 1e-3)
# final_results['im_proposals'] = final_results['im_proposals'][index, :]
# final_results['im_scores'] = final_results['im_scores'][index]
# final_results['im_labels'] = final_results['im_labels'][index]
return final_results