File size: 8,123 Bytes
26225c5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 |
from torch import nn
from copy import copy
from itertools import product
from src.utils.instance import instance_cut_pursuit
__all__ = ['InstancePartitioner']
class InstancePartitioner(nn.Module):
"""Partition a graph into instances using cut-pursuit.
More specifically, this step will group nodes together based on:
- node offset position
- node predicted classification logits
- node size
- edge affinity
NB: This operation relies on the parallel cut-pursuit algorithm:
https://gitlab.com/1a7r0ch3/parallel-cut-pursuit
Currently, this implementation is non-differentiable and runs on
CPU.
:param loss_type: str
Rules the loss applied on the node features. Accepts one of
'l2' (L2 loss on node features and probabilities),
'l2_kl' (L2 loss on node features and Kullback-Leibler
divergence on node probabilities)
:param regularization: float
Regularization parameter for the partition
:param x_weight: float
Weight used to mitigate the impact of the node position in the
partition. The larger, the less spatial coordinates matter
:param p_weight: float
Weight used to mitigate the impact of the node probabilities in
the partition. The larger, the greater the impact
:param cutoff: float
Minimum number of points in each cluster
:param parallel: bool
Whether cut-pursuit should run in parallel
:param iterations: int
Maximum number of iterations for each partition
:param trim: bool
Whether the input graph should be trimmed. See `to_trimmed()`
documentation for more details on this operation
:param discrepancy_epsilon: float
Mitigates the maximum discrepancy. More precisely:
`affinity=1 ⇒ discrepancy=1/discrepancy_epsilon`
:param temperature: float
Temperature used in the softmax when converting node logits to
probabilities
:param dampening: float
Dampening applied to the node probabilities to mitigate the
impact of near-zero probabilities in the Kullback-Leibler
divergence
:return:
"""
def __init__(
self,
loss_type='l2_kl',
regularization=10,
x_weight=1e-2,
p_weight=1,
cutoff=1,
parallel=True,
iterations=10,
trim=False,
discrepancy_epsilon=1e-4,
temperature=1,
dampening=0):
super().__init__()
self.loss_type = loss_type
self.regularization = regularization
self.x_weight = x_weight
self.p_weight = p_weight
self.cutoff = cutoff
self.parallel = parallel
self.iterations = iterations
self.trim = trim
self.discrepancy_epsilon = discrepancy_epsilon
self.temperature = temperature
self.dampening = dampening
def forward(
self,
batch,
node_x,
node_logits,
stuff_classes,
node_size,
edge_index,
edge_affinity_logits,
grid=None):
"""The forward step will compute the partition on the instance
graph, based on the node features, node logits, and edge
affinities. The partition segments will then be further merged
so that there is at most one instance of each stuff class per
batch item (ie per scene).
:param batch: Tensor of shape [num_nodes]
Batch index of each node
:param node_x: Tensor of shape [num_nodes, num_dim]
Predicted node embeddings
:param node_logits: Tensor of shape [num_nodes, num_classes]
Predicted classification logits for each node
:param stuff_classes: List or Tensor
List of 'stuff' class labels. These are used for merging
stuff segments together to ensure there is at most one
predicted instance of each 'stuff' class per batch item
:param node_size: Tensor of shape [num_nodes]
Size of each node
:param edge_index: Tensor of shape [2, num_edges]
Edges of the graph, in torch-geometric's format
:param edge_affinity_logits: Tensor of shape [num_edges]
Predicted affinity logits (ie in R+, before sigmoid) of each
edge
:param grid: Dict
A dictionary containing settings for grid-searching optimal
partition parameters
:return: obj_index: Tensor of shape [num_nodes] (or List(Dict, Tensor))
Indicates which predicted instance each node belongs to. If
a grid is passed as input, a list containing partition
settings and partition index tensors will be returned
"""
# If grid is passed, multiple partition will be computed on the
# parameter grid
if grid is not None and len(grid) > 0:
return self._grid_forward(
batch,
node_x,
node_logits,
stuff_classes,
node_size,
edge_index,
edge_affinity_logits,
grid)
# If not grid searching optimal partition parameters, simply run
# the partition with the current parameters
return instance_cut_pursuit(
batch,
node_x,
node_logits,
stuff_classes,
node_size,
edge_index,
edge_affinity_logits,
loss_type=self.loss_type,
regularization=self.regularization,
x_weight=self.x_weight,
p_weight=self.p_weight,
cutoff=self.cutoff,
parallel=self.parallel,
iterations=self.iterations,
trim=self.trim,
discrepancy_epsilon=self.discrepancy_epsilon,
temperature=self.temperature,
dampening=self.dampening)
def _grid_forward(
self,
batch,
node_x,
node_logits,
stuff_classes,
node_size,
edge_index,
edge_affinity_logits,
grid):
"""Run multiple forward calls for grid-searching optimal
settings.
"""
# If a grid dictionary was passed, make sure all keys in the
# grid are supported attributes
keys = list(grid.keys())
for k in keys:
if k not in self.__dict__:
raise ValueError(
f"'{k}' is not {self.__class__.__name__} attribute")
# Backup the current attributes
attr_bckp = copy(self.__dict__)
# Compute the grid search on the Cartesian product of the sets
# of explored values
grid_outputs = []
for values in product(*grid.values()):
# Update self attributes with grid values
for k, v in zip(keys, values):
setattr(self, k, v)
# Compute the partition
obj_index = self.forward(
batch,
node_x,
node_logits,
stuff_classes,
node_size,
edge_index,
edge_affinity_logits,
grid=None)
# Store the partition index for the current settings. The
# results are stored in a tuple whose first element is a
# dictionary of settings for self, and the second is the
# output partition index
grid_outputs.append({k: v for k, v in zip(keys, values)}, obj_index)
# Restore the initial attributes
for k, v in attr_bckp.items():
setattr(self, k, v)
return grid_outputs
def extra_repr(self) -> str:
keys = [
'regularization',
'x_weight',
'cutoff',
'parallel',
'iterations',
'trim',
'discrepancy_epsilon']
return ', '.join([f'{k}={getattr(self, k)}' for k in keys])
|