init

Author: chrischoy

.gitignore

@@ -0,0 +1,20 @@
+# Byte-compiled / optimized
+__pycache__/
+*.py[cod]
+*.o
+*.so
+
+# Text edits
+*.swp
+*.swo
+*.orig
+
+# Training
+outputs/
+
+# Profiling
+callgrind.out*
+*.dSYM
+
+# Misc
+.DS_Store

.style.yapf

@@ -0,0 +1,4 @@
+[style]
+based_on_style = chromium
+column_limit = 100
+indent_width = 2

README.md

@@ -0,0 +1,3 @@
+# Spatio-Temporal Segmentation
+
+This repository contains the accompanying code for [4D-SpatioTemporal ConvNets: Minkowski Convolutional Neural Networks, CVPR'19](https://arxiv.org/abs/1904.08755).

lib/__init__.py

@@ -0,0 +1,34 @@
+import torch
+from torch.utils.data.sampler import Sampler
+
+
+class InfSampler(Sampler):
+  """Samples elements randomly, without replacement.
+
+    Arguments:
+        data_source (Dataset): dataset to sample from
+    """
+
+  def __init__(self, data_source, shuffle=False):
+    self.data_source = data_source
+    self.shuffle = shuffle
+    self.reset_permutation()
+
+  def reset_permutation(self):
+    perm = len(self.data_source)
+    if self.shuffle:
+      perm = torch.randperm(perm)
+    self._perm = perm.tolist()
+
+  def __iter__(self):
+    return self
+
+  def __next__(self):
+    if len(self._perm) == 0:
+      self.reset_permutation()
+    return self._perm.pop()
+
+  def __len__(self):
+    return len(self.data_source)
+
+  next = __next__  # Python 2 compatibility

lib/dataloader.py

@@ -0,0 +1,85 @@
+import torch
+import torch.nn as nn
+
+from MinkowskiEngine import MinkowskiGlobalPooling, MinkowskiBroadcastAddition, MinkowskiBroadcastMultiplication
+
+
+class MinkowskiLayerNorm(nn.Module):
+
+  def __init__(self, num_features, eps=1e-5, D=-1):
+    super(MinkowskiLayerNorm, self).__init__()
+    self.num_features = num_features
+    self.eps = eps
+    self.weight = nn.Parameter(torch.ones(1, num_features))
+    self.bias = nn.Parameter(torch.zeros(1, num_features))
+
+    self.mean_in = MinkowskiGlobalPooling(dimension=D)
+    self.glob_sum = MinkowskiBroadcastAddition(dimension=D)
+    self.glob_sum2 = MinkowskiBroadcastAddition(dimension=D)
+    self.glob_mean = MinkowskiGlobalPooling(dimension=D)
+    self.glob_times = MinkowskiBroadcastMultiplication(dimension=D)
+    self.D = D
+    self.reset_parameters()
+
+  def __repr__(self):
+    s = f'(D={self.D})'
+    return self.__class__.__name__ + s
+
+  def reset_parameters(self):
+    self.weight.data.fill_(1)
+    self.bias.data.zero_()
+
+  def _check_input_dim(self, input):
+    if input.F.dim() != 2:
+      raise ValueError('expected 2D input (got {}D input)'.format(input.dim()))
+
+  def forward(self, x):
+    self._check_input_dim(x)
+    mean = self.mean_in(x).F.mean(-1, keepdim=True)
+    mean = mean + torch.zeros(mean.size(0), self.num_features).type_as(mean)
+    temp = self.glob_sum(x.F, -mean)**2
+    var = self.glob_mean(temp.data).mean(-1, keepdim=True)
+    var = var + torch.zeros(var.size(0), self.num_features).type_as(var)
+    instd = 1 / (var + self.eps).sqrt()
+
+    x = self.glob_times(self.glob_sum2(x, -mean), instd)
+    return x * self.weight + self.bias
+
+
+class MinkowskiInstanceNorm(nn.Module):
+
+  def __init__(self, num_features, eps=1e-5, D=-1):
+    super(MinkowskiInstanceNorm, self).__init__()
+    self.eps = eps
+    self.weight = nn.Parameter(torch.ones(1, num_features))
+    self.bias = nn.Parameter(torch.zeros(1, num_features))
+
+    self.mean_in = MinkowskiGlobalPooling(dimension=D)
+    self.glob_sum = MinkowskiBroadcastAddition(dimension=D)
+    self.glob_sum2 = MinkowskiBroadcastAddition(dimension=D)
+    self.glob_mean = MinkowskiGlobalPooling(dimension=D)
+    self.glob_times = MinkowskiBroadcastMultiplication(dimension=D)
+    self.D = D
+    self.reset_parameters()
+
+  def __repr__(self):
+    s = f'(pixel_dist={self.pixel_dist}, D={self.D})'
+    return self.__class__.__name__ + s
+
+  def reset_parameters(self):
+    self.weight.data.fill_(1)
+    self.bias.data.zero_()
+
+  def _check_input_dim(self, input):
+    if input.dim() != 2:
+      raise ValueError('expected 2D input (got {}D input)'.format(input.dim()))
+
+  def forward(self, x):
+    self._check_input_dim(x)
+    mean_in = self.mean_in(x)
+    temp = self.glob_sum(x, -mean_in)**2
+    var_in = self.glob_mean(temp.data)
+    instd_in = 1 / (var_in + self.eps).sqrt()
+
+    x = self.glob_times(self.glob_sum2(x, -mean_in), instd_in)
+    return x * self.weight + self.bias

lib/layers.py

@@ -0,0 +1,70 @@
+from scipy.sparse import csr_matrix
+import torch
+
+
+class SparseMM(torch.autograd.Function):
+  """
+  Sparse x dense matrix multiplication with autograd support.
+  Implementation by Soumith Chintala:
+  https://discuss.pytorch.org/t/
+  does-pytorch-support-autograd-on-sparse-matrix/6156/7
+  """
+
+  def forward(self, matrix1, matrix2):
+    self.save_for_backward(matrix1, matrix2)
+    return torch.mm(matrix1, matrix2)
+
+  def backward(self, grad_output):
+    matrix1, matrix2 = self.saved_tensors
+    grad_matrix1 = grad_matrix2 = None
+
+    if self.needs_input_grad[0]:
+        grad_matrix1 = torch.mm(grad_output, matrix2.t())
+
+    if self.needs_input_grad[1]:
+        grad_matrix2 = torch.mm(matrix1.t(), grad_output)
+
+    return grad_matrix1, grad_matrix2
+
+
+def sparse_float_tensor(values, indices, size=None):
+  """
+  Return a torch sparse matrix give values and indices (row_ind, col_ind).
+  If the size is an integer, return a square matrix with side size.
+  If the size is a torch.Size, use it to initialize the out tensor.
+  If none, the size is inferred.
+  """
+  indices = torch.stack(indices).int()
+  sargs = [indices, values.float()]
+  if size is not None:
+    # Use the provided size
+    if isinstance(size, int):
+      size = torch.Size((size, size))
+    sargs.append(size)
+  if values.is_cuda:
+    return torch.cuda.sparse.FloatTensor(*sargs)
+  else:
+    return torch.sparse.FloatTensor(*sargs)
+
+
+def diags(values, size=None):
+  values = values.view(-1)
+  n = values.nelement()
+  size = torch.Size((n, n))
+  indices = (torch.arange(0, n), torch.arange(0, n))
+  return sparse_float_tensor(values, indices, size)
+
+
+def sparse_to_csr_matrix(tensor):
+  tensor = tensor.cpu()
+  inds = tensor._indices().numpy()
+  vals = tensor._values().numpy()
+  return csr_matrix((vals, (inds[0], inds[1])), shape=[s for s in tensor.shape])
+
+
+def csr_matrix_to_sparse(mat):
+  row_ind, col_ind = mat.nonzero()
+  return sparse_float_tensor(
+      torch.from_numpy(mat.data),
+      (torch.from_numpy(row_ind), torch.from_numpy(col_ind)),
+      size=torch.Size(mat.shape))

lib/math_functions.py

@@ -0,0 +1,78 @@
+import logging
+
+from torch.optim import SGD, Adam
+from torch.optim.lr_scheduler import LambdaLR, StepLR
+
+
+class LambdaStepLR(LambdaLR):
+
+  def __init__(self, optimizer, lr_lambda, last_step=-1):
+    super(LambdaStepLR, self).__init__(optimizer, lr_lambda, last_step)
+
+  @property
+  def last_step(self):
+    """Use last_epoch for the step counter"""
+    return self.last_epoch
+
+  @last_step.setter
+  def last_step(self, v):
+    self.last_epoch = v
+
+
+class PolyLR(LambdaStepLR):
+  """DeepLab learning rate policy"""
+
+  def __init__(self, optimizer, max_iter, power=0.9, last_step=-1):
+    super(PolyLR, self).__init__(optimizer, lambda s: (1 - s / (max_iter + 1))**power, last_step)
+
+
+class SquaredLR(LambdaStepLR):
+  """ Used for SGD Lars"""
+
+  def __init__(self, optimizer, max_iter, last_step=-1):
+    super(SquaredLR, self).__init__(optimizer, lambda s: (1 - s / (max_iter + 1))**2, last_step)
+
+
+class ExpLR(LambdaStepLR):
+
+  def __init__(self, optimizer, step_size, gamma=0.9, last_step=-1):
+    # (0.9 ** 21.854) = 0.1, (0.95 ** 44.8906) = 0.1
+    # To get 0.1 every N using gamma 0.9, N * log(0.9)/log(0.1) = 0.04575749 N
+    # To get 0.1 every N using gamma g, g ** N = 0.1 -> N * log(g) = log(0.1) -> g = np.exp(log(0.1) / N)
+    super(ExpLR, self).__init__(optimizer, lambda s: gamma**(s / step_size), last_step)
+
+
+def initialize_optimizer(params, config):
+  assert config.optimizer in ['SGD', 'Adagrad', 'Adam', 'RMSProp', 'Rprop', 'SGDLars']
+
+  if config.optimizer == 'SGD':
+    return SGD(
+        params,
+        lr=config.lr,
+        momentum=config.sgd_momentum,
+        dampening=config.sgd_dampening,
+        weight_decay=config.weight_decay)
+  elif config.optimizer == 'Adam':
+    return Adam(
+        params,
+        lr=config.lr,
+        betas=(config.adam_beta1, config.adam_beta2),
+        weight_decay=config.weight_decay)
+  else:
+    logging.error('Optimizer type not supported')
+    raise ValueError('Optimizer type not supported')
+
+
+def initialize_scheduler(optimizer, config, last_step=-1):
+  if config.scheduler == 'StepLR':
+    return StepLR(
+        optimizer, step_size=config.step_size, gamma=config.step_gamma, last_epoch=last_step)
+  elif config.scheduler == 'PolyLR':
+    return PolyLR(optimizer, max_iter=config.max_iter, power=config.poly_power, last_step=last_step)
+  elif config.scheduler == 'SquaredLR':
+    return SquaredLR(optimizer, max_iter=config.max_iter, last_step=last_step)
+  elif config.scheduler == 'ExpLR':
+    return ExpLR(
+        optimizer, step_size=config.exp_step_size, gamma=config.exp_gamma, last_step=last_step)
+  else:
+    logging.error('Scheduler not supported')