models
dreem.models
¶
Model architectures and layers.
Modules:
| Name | Description |
|---|---|
attention_head |
Module containing different components of multi-head attention heads. |
embedding |
Module containing different position and temporal embeddings. |
global_tracking_transformer |
Module containing GTR model used for training. |
gtr_runner |
Module containing training, validation and inference logic. |
mlp |
Multi-Layer Perceptron (MLP) module. |
model_utils |
Module containing model helper functions. |
transformer |
DETR Transformer class. |
visual_encoder |
Module for different visual feature extractors. |
Classes:
| Name | Description |
|---|---|
DescriptorVisualEncoder |
Visual Encoder based on image descriptors. |
Embedding |
Class that wraps around different embedding types. |
FourierPositionalEmbeddings |
Fourier positional embeddings. |
GTRRunner |
A lightning wrapper around GTR model. |
GlobalTrackingTransformer |
Modular GTR model composed of visual encoder + transformer used for tracking. |
Transformer |
Transformer class. |
VisualEncoder |
Class wrapping around a visual feature extractor backbone. |
Functions:
| Name | Description |
|---|---|
create_visual_encoder |
Create a visual encoder based on the specified type. |
register_encoder |
Register a new encoder type. |
DescriptorVisualEncoder
¶
Bases: Module
Visual Encoder based on image descriptors.
Methods:
| Name | Description |
|---|---|
__init__ |
Initialize Descriptor Visual Encoder. |
compute_hu_moments |
Compute Hu moments. |
compute_inertia_tensor |
Compute inertia tensor. |
forward |
Forward pass of feature extractor to get feature vector. |
Source code in dreem/models/visual_encoder.py
class DescriptorVisualEncoder(torch.nn.Module):
"""Visual Encoder based on image descriptors."""
def __init__(self, use_hu_moments: bool = False, **kwargs):
"""Initialize Descriptor Visual Encoder.
Args:
use_hu_moments: Whether to use Hu moments.
**kwargs: Additional keyword arguments (unused but accepted for compatibility).
"""
super().__init__()
self.use_hu_moments = use_hu_moments
def compute_hu_moments(self, img):
"""Compute Hu moments."""
mu = measure.moments_central(img)
nu = measure.moments_normalized(mu)
hu = measure.moments_hu(nu)
# log transform hu moments for scale differences; switched off; numerically unstable
# hu_log = -np.sign(hu) * np.log(np.abs(hu))
return hu
def compute_inertia_tensor(self, img):
"""Compute inertia tensor."""
return measure.inertia_tensor(img)
@torch.no_grad()
def forward(self, img: torch.Tensor) -> torch.Tensor:
"""Forward pass of feature extractor to get feature vector."""
descriptors = []
for im in img:
im = im[0].cpu().numpy()
inertia_tensor = self.compute_inertia_tensor(im)
mean_intensity = im.mean()
if self.use_hu_moments:
hu_moments = self.compute_hu_moments(im)
# Flatten inertia tensor
inertia_tensor_flat = inertia_tensor.flatten()
# Combine all features into a single descriptor
descriptor = np.concatenate(
[
inertia_tensor_flat,
[mean_intensity],
hu_moments if self.use_hu_moments else [],
]
)
descriptors.append(torch.tensor(descriptor, dtype=torch.float32))
return torch.stack(descriptors)
__init__(use_hu_moments=False, **kwargs)
¶
Initialize Descriptor Visual Encoder.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
use_hu_moments
|
bool
|
Whether to use Hu moments. |
False
|
**kwargs
|
Additional keyword arguments (unused but accepted for compatibility). |
{}
|
Source code in dreem/models/visual_encoder.py
compute_hu_moments(img)
¶
Compute Hu moments.
Source code in dreem/models/visual_encoder.py
compute_inertia_tensor(img)
¶
forward(img)
¶
Forward pass of feature extractor to get feature vector.
Source code in dreem/models/visual_encoder.py
@torch.no_grad()
def forward(self, img: torch.Tensor) -> torch.Tensor:
"""Forward pass of feature extractor to get feature vector."""
descriptors = []
for im in img:
im = im[0].cpu().numpy()
inertia_tensor = self.compute_inertia_tensor(im)
mean_intensity = im.mean()
if self.use_hu_moments:
hu_moments = self.compute_hu_moments(im)
# Flatten inertia tensor
inertia_tensor_flat = inertia_tensor.flatten()
# Combine all features into a single descriptor
descriptor = np.concatenate(
[
inertia_tensor_flat,
[mean_intensity],
hu_moments if self.use_hu_moments else [],
]
)
descriptors.append(torch.tensor(descriptor, dtype=torch.float32))
return torch.stack(descriptors)
Embedding
¶
Bases: Module
Class that wraps around different embedding types.
Used for both learned and fixed embeddings.
Methods:
| Name | Description |
|---|---|
__init__ |
Initialize embeddings. |
forward |
Get the sequence positional embeddings. |
Source code in dreem/models/embedding.py
class Embedding(torch.nn.Module):
"""Class that wraps around different embedding types.
Used for both learned and fixed embeddings.
"""
EMB_TYPES = {
"temp": {},
"pos": {"over_boxes"},
"off": {},
None: {},
} # dict of valid args:keyword params
EMB_MODES = {
"fixed": {"temperature", "scale", "normalize"},
"learned": {"emb_num"},
"off": {},
} # dict of valid args:keyword params
def __init__(
self,
emb_type: str,
mode: str,
features: int,
n_points: int = 1,
emb_num: int = 16,
over_boxes: bool = True,
temperature: int = 10000,
normalize: bool = False,
scale: float | None = None,
mlp_cfg: dict | None = None,
):
"""Initialize embeddings.
Args:
emb_type: The type of embedding to compute.
Must be one of `{"temp", "pos", "off"}`
mode: The mode or function used to map positions to vector embeddings.
Must be one of `{"fixed", "learned", "off"}`
features: The embedding dimensions. Must match the dimension of the
input vectors for the transformer model.
n_points: the number of points that will be embedded.
emb_num: the number of embeddings in the `self.lookup` table
(Only used in learned embeddings).
over_boxes: Whether to compute the position embedding for each bbox
coordinate (y1x1y2x2) or the centroid + bbox size (yxwh).
temperature: the temperature constant to be used when computing
the sinusoidal position embedding
normalize: whether or not to normalize the positions
(Only used in fixed embeddings).
scale: factor by which to scale the positions after normalizing
(Only used in fixed embeddings).
mlp_cfg: A dictionary of mlp hyperparameters for projecting
embedding to correct space.
Example: {"hidden_dims": 256, "num_layers":3, "dropout": 0.3}
"""
self._check_init_args(emb_type, mode)
super().__init__()
self.emb_type = emb_type
self.mode = mode
self.features = features
self.emb_num = emb_num
self.over_boxes = over_boxes
self.temperature = temperature
self.normalize = normalize
self.scale = scale
self.n_points = n_points
if self.normalize and self.scale is None:
self.scale = 2 * math.pi
if self.emb_type == "pos" and mlp_cfg is not None and mlp_cfg["num_layers"] > 0:
if self.mode == "fixed":
self.mlp = MLP(
input_dim=n_points * self.features,
output_dim=self.features,
**mlp_cfg,
)
else:
in_dim = (self.features // (4 * n_points)) * (4 * n_points)
self.mlp = MLP(
input_dim=in_dim,
output_dim=self.features,
**mlp_cfg,
)
else:
self.mlp = torch.nn.Identity()
self._emb_func = lambda tensor: torch.zeros(
(tensor.shape[0], self.features), dtype=tensor.dtype, device=tensor.device
) # turn off embedding by returning zeros
self.lookup = None
if self.mode == "learned":
if self.emb_type == "pos":
self.lookup = torch.nn.Embedding(
self.emb_num * 4 * self.n_points, self.features // (4 * n_points)
)
self._emb_func = self._learned_pos_embedding
elif self.emb_type == "temp":
self.lookup = torch.nn.Embedding(self.emb_num, self.features)
self._emb_func = self._learned_temp_embedding
elif self.mode == "fixed":
if self.emb_type == "pos":
self._emb_func = self._sine_box_embedding
elif self.emb_type == "temp":
self._emb_func = self._sine_temp_embedding
def _check_init_args(self, emb_type: str, mode: str):
"""Check whether the correct arguments were passed to initialization.
Args:
emb_type: The type of embedding to compute. Must be one of `{"temp", "pos", ""}`
mode: The mode or function used to map positions to vector embeddings.
Must be one of `{"fixed", "learned"}`
Raises:
ValueError:
* if the incorrect `emb_type` or `mode` string are passed
NotImplementedError: if `emb_type` is `temp` and `mode` is `fixed`.
"""
if emb_type.lower() not in self.EMB_TYPES:
raise ValueError(
f"Embedding `emb_type` must be one of {self.EMB_TYPES} not {emb_type}"
)
if mode.lower() not in self.EMB_MODES:
raise ValueError(
f"Embedding `mode` must be one of {self.EMB_MODES} not {mode}"
)
def forward(self, seq_positions: torch.Tensor) -> torch.Tensor:
"""Get the sequence positional embeddings.
Args:
seq_positions:
* An (`N`, 1) tensor where seq_positions[i] represents the temporal position of instance_i in the sequence.
* An (`N`, n_anchors x 4) tensor where seq_positions[i, j, :] represents the [y1, x1, y2, x2] spatial locations of jth point of instance_i in the sequence.
Returns:
An `N` x `self.features` tensor representing the corresponding spatial or temporal embedding.
"""
emb = self._emb_func(seq_positions)
if emb.shape[-1] != self.features:
raise RuntimeError(
(
f"Output embedding dimension is {emb.shape[-1]} but requested {self.features} dimensions! \n"
f"hint: Try turning the MLP on by passing `mlp_cfg` to the constructor to project to the correct embedding dimensions."
)
)
return emb
def _torch_int_div(
self, tensor1: torch.Tensor, tensor2: torch.Tensor
) -> torch.Tensor:
"""Perform integer division of two tensors.
Args:
tensor1: dividend tensor.
tensor2: divisor tensor.
Returns:
torch.Tensor, resulting tensor.
"""
return torch.div(tensor1, tensor2, rounding_mode="floor")
def _sine_box_embedding(self, boxes: torch.Tensor) -> torch.Tensor:
"""Compute sine positional embeddings for boxes using given parameters.
Args:
boxes: the input boxes of shape N, n_anchors, 4 or B, N, n_anchors, 4
where the last dimension is the bbox coords in [y1, x1, y2, x2].
(Note currently `B=batch_size=1`).
Returns:
torch.Tensor, the sine positional embeddings
(embedding[:, 4i] = sin(x)
embedding[:, 4i+1] = cos(x)
embedding[:, 4i+2] = sin(y)
embedding[:, 4i+3] = cos(y)
)
"""
if self.scale is not None and self.normalize is False:
raise ValueError("normalize should be True if scale is passed")
if len(boxes.size()) == 3:
boxes = boxes.unsqueeze(0)
if self.normalize:
boxes = boxes / (boxes[:, :, -1:] + 1e-6) * self.scale
dim_t = torch.arange(self.features // 4, dtype=torch.float32)
dim_t = self.temperature ** (
2 * self._torch_int_div(dim_t, 2) / (self.features // 4)
)
# (b, n_t, n_anchors, 4, D//4)
pos_emb = boxes[:, :, :, :, None] / dim_t.to(boxes.device)
pos_emb = torch.stack(
(pos_emb[:, :, :, :, 0::2].sin(), pos_emb[:, :, :, :, 1::2].cos()), dim=4
)
pos_emb = pos_emb.flatten(2).squeeze(0) # (N_t, n_anchors * D)
pos_emb = self.mlp(pos_emb)
pos_emb = pos_emb.view(boxes.shape[1], self.features)
return pos_emb
def _sine_temp_embedding(self, times: torch.Tensor) -> torch.Tensor:
"""Compute fixed sine temporal embeddings.
Args:
times: the input times of shape (N,) or (N,1) where N = (sum(instances_per_frame))
which is the frame index of the instance relative
to the batch size
(e.g. `torch.tensor([0, 0, ..., 0, 1, 1, ..., 1, 2, 2, ..., 2,..., B, B, ...B])`).
Returns:
an n_instances x D embedding representing the temporal embedding.
"""
T = times.int().max().item() + 1
d = self.features
n = self.temperature
positions = torch.arange(0, T).unsqueeze(1)
temp_lookup = torch.zeros(T, d, device=times.device)
denominators = torch.pow(
n, 2 * torch.arange(0, d // 2) / d
) # 10000^(2i/d_model), i is the index of embedding
temp_lookup[:, 0::2] = torch.sin(
positions / denominators
) # sin(pos/10000^(2i/d_model))
temp_lookup[:, 1::2] = torch.cos(
positions / denominators
) # cos(pos/10000^(2i/d_model))
temp_emb = temp_lookup[times.int()]
return temp_emb # .view(len(times), self.features)
def _learned_pos_embedding(self, boxes: torch.Tensor) -> torch.Tensor:
"""Compute learned positional embeddings for boxes using given parameters.
Args:
boxes: the input boxes of shape N x 4 or B x N x 4
where the last dimension is the bbox coords in [y1, x1, y2, x2].
(Note currently `B=batch_size=1`).
Returns:
torch.Tensor, the learned positional embeddings.
"""
pos_lookup = self.lookup
N, n_anchors, _ = boxes.shape
boxes = boxes.view(N, n_anchors, 4)
if self.over_boxes:
xywh = boxes
else:
xywh = torch.cat(
[
(boxes[:, :, 2:] + boxes[:, :, :2]) / 2,
(boxes[:, :, 2:] - boxes[:, :, :2]),
],
dim=1,
)
left_ind, right_ind, left_weight, right_weight = self._compute_weights(xywh)
f = pos_lookup.weight.shape[1] # self.features // 4
try:
pos_emb_table = pos_lookup.weight.view(
self.emb_num, n_anchors, 4, f
) # T x 4 x (D * 4)
except RuntimeError as e:
logger.exception(
f"Hint: `n_points` ({self.n_points}) may be set incorrectly!"
)
logger.exception(e)
raise (e)
left_emb = pos_emb_table.gather(
0,
left_ind[:, :, :, None].to(pos_emb_table.device).expand(N, n_anchors, 4, f),
) # N x 4 x d
right_emb = pos_emb_table.gather(
0,
right_ind[:, :, :, None]
.to(pos_emb_table.device)
.expand(N, n_anchors, 4, f),
) # N x 4 x d
pos_emb = left_weight[:, :, :, None] * right_emb.to(
left_weight.device
) + right_weight[:, :, :, None] * left_emb.to(right_weight.device)
pos_emb = pos_emb.flatten(1)
pos_emb = self.mlp(pos_emb)
return pos_emb.view(N, self.features)
def _learned_temp_embedding(self, times: torch.Tensor) -> torch.Tensor:
"""Compute learned temporal embeddings for times using given parameters.
Args:
times: the input times of shape (N,) or (N,1) where N = (sum(instances_per_frame))
which is the frame index of the instance relative
to the batch size
(e.g. `torch.tensor([0, 0, ..., 0, 1, 1, ..., 1, 2, 2, ..., 2,..., B, B, ...B])`).
Returns:
torch.Tensor, the learned temporal embeddings.
"""
temp_lookup = self.lookup
N = times.shape[0]
left_ind, right_ind, left_weight, right_weight = self._compute_weights(times)
left_emb = temp_lookup.weight[
left_ind.to(temp_lookup.weight.device)
] # T x D --> N x D
right_emb = temp_lookup.weight[right_ind.to(temp_lookup.weight.device)]
temp_emb = left_weight[:, None] * right_emb.to(
left_weight.device
) + right_weight[:, None] * left_emb.to(right_weight.device)
return temp_emb.view(N, self.features)
def _compute_weights(self, data: torch.Tensor) -> tuple[torch.Tensor, ...]:
"""Compute left and right learned embedding weights.
Args:
data: the input data (e.g boxes or times).
Returns:
A torch.Tensor for each of the left/right indices and weights, respectively
"""
data = data * self.emb_num
left_ind = data.clamp(min=0, max=self.emb_num - 1).long() # N x 4
right_ind = (left_ind + 1).clamp(min=0, max=self.emb_num - 1).long() # N x 4
left_weight = data - left_ind.float() # N x 4
right_weight = 1.0 - left_weight
return left_ind, right_ind, left_weight, right_weight
__init__(emb_type, mode, features, n_points=1, emb_num=16, over_boxes=True, temperature=10000, normalize=False, scale=None, mlp_cfg=None)
¶
Initialize embeddings.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
emb_type
|
str
|
The type of embedding to compute.
Must be one of |
required |
mode
|
str
|
The mode or function used to map positions to vector embeddings.
Must be one of |
required |
features
|
int
|
The embedding dimensions. Must match the dimension of the input vectors for the transformer model. |
required |
n_points
|
int
|
the number of points that will be embedded. |
1
|
emb_num
|
int
|
the number of embeddings in the |
16
|
over_boxes
|
bool
|
Whether to compute the position embedding for each bbox coordinate (y1x1y2x2) or the centroid + bbox size (yxwh). |
True
|
temperature
|
int
|
the temperature constant to be used when computing the sinusoidal position embedding |
10000
|
normalize
|
bool
|
whether or not to normalize the positions (Only used in fixed embeddings). |
False
|
scale
|
float | None
|
factor by which to scale the positions after normalizing (Only used in fixed embeddings). |
None
|
mlp_cfg
|
dict | None
|
A dictionary of mlp hyperparameters for projecting embedding to correct space. Example: {"hidden_dims": 256, "num_layers":3, "dropout": 0.3} |
None
|
Source code in dreem/models/embedding.py
def __init__(
self,
emb_type: str,
mode: str,
features: int,
n_points: int = 1,
emb_num: int = 16,
over_boxes: bool = True,
temperature: int = 10000,
normalize: bool = False,
scale: float | None = None,
mlp_cfg: dict | None = None,
):
"""Initialize embeddings.
Args:
emb_type: The type of embedding to compute.
Must be one of `{"temp", "pos", "off"}`
mode: The mode or function used to map positions to vector embeddings.
Must be one of `{"fixed", "learned", "off"}`
features: The embedding dimensions. Must match the dimension of the
input vectors for the transformer model.
n_points: the number of points that will be embedded.
emb_num: the number of embeddings in the `self.lookup` table
(Only used in learned embeddings).
over_boxes: Whether to compute the position embedding for each bbox
coordinate (y1x1y2x2) or the centroid + bbox size (yxwh).
temperature: the temperature constant to be used when computing
the sinusoidal position embedding
normalize: whether or not to normalize the positions
(Only used in fixed embeddings).
scale: factor by which to scale the positions after normalizing
(Only used in fixed embeddings).
mlp_cfg: A dictionary of mlp hyperparameters for projecting
embedding to correct space.
Example: {"hidden_dims": 256, "num_layers":3, "dropout": 0.3}
"""
self._check_init_args(emb_type, mode)
super().__init__()
self.emb_type = emb_type
self.mode = mode
self.features = features
self.emb_num = emb_num
self.over_boxes = over_boxes
self.temperature = temperature
self.normalize = normalize
self.scale = scale
self.n_points = n_points
if self.normalize and self.scale is None:
self.scale = 2 * math.pi
if self.emb_type == "pos" and mlp_cfg is not None and mlp_cfg["num_layers"] > 0:
if self.mode == "fixed":
self.mlp = MLP(
input_dim=n_points * self.features,
output_dim=self.features,
**mlp_cfg,
)
else:
in_dim = (self.features // (4 * n_points)) * (4 * n_points)
self.mlp = MLP(
input_dim=in_dim,
output_dim=self.features,
**mlp_cfg,
)
else:
self.mlp = torch.nn.Identity()
self._emb_func = lambda tensor: torch.zeros(
(tensor.shape[0], self.features), dtype=tensor.dtype, device=tensor.device
) # turn off embedding by returning zeros
self.lookup = None
if self.mode == "learned":
if self.emb_type == "pos":
self.lookup = torch.nn.Embedding(
self.emb_num * 4 * self.n_points, self.features // (4 * n_points)
)
self._emb_func = self._learned_pos_embedding
elif self.emb_type == "temp":
self.lookup = torch.nn.Embedding(self.emb_num, self.features)
self._emb_func = self._learned_temp_embedding
elif self.mode == "fixed":
if self.emb_type == "pos":
self._emb_func = self._sine_box_embedding
elif self.emb_type == "temp":
self._emb_func = self._sine_temp_embedding
forward(seq_positions)
¶
Get the sequence positional embeddings.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
seq_positions
|
Tensor
|
|
required |
Returns:
| Type | Description |
|---|---|
Tensor
|
An |
Source code in dreem/models/embedding.py
def forward(self, seq_positions: torch.Tensor) -> torch.Tensor:
"""Get the sequence positional embeddings.
Args:
seq_positions:
* An (`N`, 1) tensor where seq_positions[i] represents the temporal position of instance_i in the sequence.
* An (`N`, n_anchors x 4) tensor where seq_positions[i, j, :] represents the [y1, x1, y2, x2] spatial locations of jth point of instance_i in the sequence.
Returns:
An `N` x `self.features` tensor representing the corresponding spatial or temporal embedding.
"""
emb = self._emb_func(seq_positions)
if emb.shape[-1] != self.features:
raise RuntimeError(
(
f"Output embedding dimension is {emb.shape[-1]} but requested {self.features} dimensions! \n"
f"hint: Try turning the MLP on by passing `mlp_cfg` to the constructor to project to the correct embedding dimensions."
)
)
return emb
FourierPositionalEmbeddings
¶
Bases: Module
Fourier positional embeddings.
Methods:
| Name | Description |
|---|---|
__init__ |
Initialize Fourier positional embeddings. |
forward |
Compute learnable fourier coefficients for each spatial/temporal position. |
Source code in dreem/models/embedding.py
class FourierPositionalEmbeddings(torch.nn.Module):
"""Fourier positional embeddings."""
def __init__(
self,
n_components: int,
d_model: int,
):
"""Initialize Fourier positional embeddings.
Args:
n_components: Number of frequencies for each dimension.
d_model: Model dimension.
"""
super().__init__()
self.d_model = d_model
self.n_components = n_components
self.freq = torch.nn.Parameter(
_pos_embed_fourier1d_init(self.d_model, n_components)
)
def forward(self, seq_positions: torch.Tensor):
"""Compute learnable fourier coefficients for each spatial/temporal position.
Args:
seq_positions: tensor of shape (num_queries,)
Returns:
tensor of shape (num_queries, embed_dim)
"""
freq = self.freq.to(seq_positions.device)
# seq_positions is of shape (num_queries,) but needs to be (1,num_queries,1)
embed = torch.cat(
(
torch.sin(
0.5 * math.pi * seq_positions.unsqueeze(-1).unsqueeze(0) * freq
),
torch.cos(
0.5 * math.pi * seq_positions.unsqueeze(-1).unsqueeze(0) * freq
),
),
axis=-1,
) / math.sqrt(len(freq)) # (B,N,2*n_components)
if self.d_model % self.n_components != 0:
raise ValueError(
f"d_model ({self.d_model}) must be divisible by number of Fourier components n_components ({self.n_components})"
)
# tile until shape is (B,N,embed_dim) to multiply with input queries/keys
embed = embed.repeat(
1, 1, self.d_model // (2 * self.n_components)
) # 2*n_components to account for sin/cos
return embed
__init__(n_components, d_model)
¶
Initialize Fourier positional embeddings.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
n_components
|
int
|
Number of frequencies for each dimension. |
required |
d_model
|
int
|
Model dimension. |
required |
Source code in dreem/models/embedding.py
def __init__(
self,
n_components: int,
d_model: int,
):
"""Initialize Fourier positional embeddings.
Args:
n_components: Number of frequencies for each dimension.
d_model: Model dimension.
"""
super().__init__()
self.d_model = d_model
self.n_components = n_components
self.freq = torch.nn.Parameter(
_pos_embed_fourier1d_init(self.d_model, n_components)
)
forward(seq_positions)
¶
Compute learnable fourier coefficients for each spatial/temporal position.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
seq_positions
|
Tensor
|
tensor of shape (num_queries,) |
required |
Returns:
| Type | Description |
|---|---|
|
tensor of shape (num_queries, embed_dim) |
Source code in dreem/models/embedding.py
def forward(self, seq_positions: torch.Tensor):
"""Compute learnable fourier coefficients for each spatial/temporal position.
Args:
seq_positions: tensor of shape (num_queries,)
Returns:
tensor of shape (num_queries, embed_dim)
"""
freq = self.freq.to(seq_positions.device)
# seq_positions is of shape (num_queries,) but needs to be (1,num_queries,1)
embed = torch.cat(
(
torch.sin(
0.5 * math.pi * seq_positions.unsqueeze(-1).unsqueeze(0) * freq
),
torch.cos(
0.5 * math.pi * seq_positions.unsqueeze(-1).unsqueeze(0) * freq
),
),
axis=-1,
) / math.sqrt(len(freq)) # (B,N,2*n_components)
if self.d_model % self.n_components != 0:
raise ValueError(
f"d_model ({self.d_model}) must be divisible by number of Fourier components n_components ({self.n_components})"
)
# tile until shape is (B,N,embed_dim) to multiply with input queries/keys
embed = embed.repeat(
1, 1, self.d_model // (2 * self.n_components)
) # 2*n_components to account for sin/cos
return embed
GTRRunner
¶
Bases: LightningModule
A lightning wrapper around GTR model.
Used for training, validation and inference.
Methods:
| Name | Description |
|---|---|
__init__ |
Initialize a lightning module for GTR. |
configure_optimizers |
Get optimizers and schedulers for training. |
forward |
Execute forward pass of the lightning module. |
log_metrics |
Log metrics computed during evaluation. |
on_test_end |
Run inference and metrics pipeline to compute metrics for test set. |
on_validation_epoch_end |
Execute hook for validation end. |
predict_step |
Run inference for model. |
test_step |
Execute single test step for model. |
training_step |
Execute single training step for model. |
validation_step |
Execute single val step for model. |
Source code in dreem/models/gtr_runner.py
class GTRRunner(LightningModule):
"""A lightning wrapper around GTR model.
Used for training, validation and inference.
"""
DEFAULT_METRICS = {
"train": [],
"val": [],
"test": ["num_switches", "global_tracking_accuracy"],
}
DEFAULT_TRACKING = {
"train": False,
"val": False,
"test": True,
}
DEFAULT_SAVE = {"train": False, "val": False, "test": False}
def __init__(
self,
model_cfg: dict | None = None,
tracker_cfg: dict | None = None,
loss_cfg: dict | None = None,
optimizer_cfg: dict | None = None,
scheduler_cfg: dict | None = None,
metrics: dict[str, list[str]] | None = None,
persistent_tracking: dict[str, bool] | None = None,
test_save_path: str = "./test_results.h5",
):
"""Initialize a lightning module for GTR.
Args:
model_cfg: hyperparameters for GlobalTrackingTransformer
tracker_cfg: The parameters used for the tracker post-processing
loss_cfg: hyperparameters for AssoLoss
optimizer_cfg: hyper parameters used for optimizer.
Only used to overwrite `configure_optimizer`
scheduler_cfg: hyperparameters for lr_scheduler used to overwrite `configure_optimizer
metrics: a dict containing the metrics to be computed during train, val, and test.
persistent_tracking: a dict containing whether to use persistent tracking during train, val and test inference.
test_save_path: path to a directory to save the eval and tracking results to
"""
super().__init__()
self.save_hyperparameters()
self.model_cfg = model_cfg if model_cfg else {}
self.loss_cfg = loss_cfg if loss_cfg else {}
self.tracker_cfg = tracker_cfg if tracker_cfg else {}
self.model = GlobalTrackingTransformer(**self.model_cfg)
self.loss = AssoLoss(**self.loss_cfg)
if self.tracker_cfg.get("tracker_type", "standard") == "batch":
from dreem.inference.batch_tracker import BatchTracker
self.tracker = BatchTracker(**self.tracker_cfg)
else:
from dreem.inference.tracker import Tracker
self.tracker = Tracker(**self.tracker_cfg)
self.optimizer_cfg = optimizer_cfg
self.scheduler_cfg = scheduler_cfg
self.metrics = metrics if metrics is not None else self.DEFAULT_METRICS
self.persistent_tracking = (
persistent_tracking
if persistent_tracking is not None
else self.DEFAULT_TRACKING
)
self.test_results = {"preds": [], "save_path": test_save_path}
def forward(
self,
ref_instances: list["Instance"],
query_instances: list["Instance"] | None = None,
) -> list["AssociationMatrix"]:
"""Execute forward pass of the lightning module.
Args:
ref_instances: a list of `Instance` objects containing crops and other data needed for transformer model
query_instances: a list of `Instance` objects used as queries in the decoder. Mostly used for inference.
Returns:
An association matrix between objects
"""
asso_preds = self.model(ref_instances, query_instances)
return asso_preds
def training_step(
self, train_batch: list[list["Frame"]], batch_idx: int
) -> dict[str, float]:
"""Execute single training step for model.
Args:
train_batch: A single batch from the dataset which is a list of `Frame` objects
with length `clip_length` containing Instances and other metadata.
batch_idx: the batch number used by lightning
Returns:
A dict containing the train loss plus any other metrics specified
"""
result = self._shared_eval_step(train_batch[0], mode="train")
self.log_metrics(result, len(train_batch[0]), "train")
return result
def validation_step(
self, val_batch: list[list["Frame"]], batch_idx: int
) -> dict[str, float]:
"""Execute single val step for model.
Args:
val_batch: A single batch from the dataset which is a list of `Frame` objects
with length `clip_length` containing Instances and other metadata.
batch_idx: the batch number used by lightning
Returns:
A dict containing the val loss plus any other metrics specified
"""
result = self._shared_eval_step(val_batch[0], mode="val")
self.log_metrics(result, len(val_batch[0]), "val")
return result
def test_step(
self, test_batch: list[list["Frame"]], batch_idx: int
) -> dict[str, float]:
"""Execute single test step for model.
Args:
test_batch: A single batch from the dataset which is a list of `Frame` objects
with length `clip_length` containing Instances and other metadata.
batch_idx: the batch number used by lightning
Returns:
A dict containing the val loss plus any other metrics specified
"""
result = self._shared_eval_step(test_batch[0], mode="test")
self.log_metrics(result, len(test_batch[0]), "test")
return result
def predict_step(self, batch: list[list["Frame"]], batch_idx: int) -> list["Frame"]:
"""Run inference for model.
Computes association + assignment.
Args:
batch: A single batch from the dataset which is a list of `Frame` objects
with length `clip_length` containing Instances and other metadata.
batch_idx: the batch number used by lightning
Returns:
A list of dicts where each dict is a frame containing the predicted track ids
"""
frames_pred = self.tracker(self.model, batch[0])
return frames_pred
def _shared_eval_step(self, frames: list["Frame"], mode: str) -> dict[str, float]:
"""Run evaluation used by train, test, and val steps.
Args:
frames: A list of dicts where each dict is a frame containing gt data
mode: which metrics to compute and whether to use persistent tracking or not
Returns:
a dict containing the loss and any other metrics specified by `eval_metrics`
"""
try:
instances = [instance for frame in frames for instance in frame.instances]
if len(instances) == 0:
return None
# eval_metrics = self.metrics[mode] # Currently unused but available for future metric computation
logits = self(instances)
logits = [asso.matrix for asso in logits]
loss = self.loss(logits, frames)
return_metrics = {"loss": loss}
if mode == "test":
self.tracker.persistent_tracking = True
frames_pred = self.tracker(self.model, frames)
self.test_results["preds"].extend(
[frame.to("cpu") for frame in frames_pred]
)
return_metrics["batch_size"] = len(frames)
except Exception as e:
logger.exception(
f"Failed on frame {frames[0].frame_id} of video {frames[0].video_id}"
)
logger.exception(e)
raise (e)
return return_metrics
def configure_optimizers(self) -> dict:
"""Get optimizers and schedulers for training.
Is overridden by config but defaults to Adam + ReduceLROnPlateau.
Returns:
an optimizer config dict containing the optimizer, scheduler, and scheduler params
"""
# todo: init from config
if self.optimizer_cfg is None:
optimizer = torch.optim.Adam(self.parameters(), lr=1e-4, betas=(0.9, 0.999))
else:
optimizer = init_optimizer(self.parameters(), self.optimizer_cfg)
if self.scheduler_cfg is None:
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
optimizer, "min", 0.5, 10
)
else:
scheduler = init_scheduler(optimizer, self.scheduler_cfg)
return {
"optimizer": optimizer,
"lr_scheduler": {
"scheduler": scheduler,
"monitor": "val_loss",
"interval": "epoch",
"frequency": 1,
},
}
def log_metrics(self, result: dict, batch_size: int, mode: str) -> None:
"""Log metrics computed during evaluation.
Args:
result: A dict containing metrics to be logged.
batch_size: the size of the batch used to compute the metrics
mode: One of {'train', 'test' or 'val'}. Used as prefix while logging.
"""
if result:
batch_size = result.pop("batch_size")
for metric, val in result.items():
if isinstance(val, torch.Tensor):
val = val.item()
self.log(f"{mode}_{metric}", val, batch_size=batch_size)
def on_validation_epoch_end(self):
"""Execute hook for validation end.
Currently, we simply clear the gpu cache and do garbage collection.
"""
gc.collect()
torch.cuda.empty_cache()
def on_test_end(self):
"""Run inference and metrics pipeline to compute metrics for test set.
Args:
test_results: dict containing predictions and metrics to be filled out in metrics.evaluate
metrics: list of metrics to compute
"""
# input validation
metrics_to_compute = self.metrics[
"test"
] # list of metrics to compute, or "all"
if metrics_to_compute == "all":
metrics_to_compute = ["motmetrics", "global_tracking_accuracy"]
if isinstance(metrics_to_compute, str):
metrics_to_compute = [metrics_to_compute]
for metric in metrics_to_compute:
if metric not in ["motmetrics", "global_tracking_accuracy"]:
raise ValueError(
f"Metric {metric} not supported. Please select from 'motmetrics' or 'global_tracking_accuracy'"
)
preds = self.test_results["preds"]
# results is a dict with key being the metric name, and value being the metric value computed
results = metrics.evaluate(preds, metrics_to_compute)
# save metrics and frame metadata to hdf5
# Get the video name from the first frame
vid_name = Path(preds[0].vid_name).stem
# save the results to an hdf5 file
fname = os.path.join(
self.test_results["save_path"], f"{vid_name}.dreem_metrics.h5"
)
logger.info(f"Saving metrics to {fname}")
# Check if the h5 file exists and add a suffix to prevent name collision
suffix_counter = 0
original_fname = fname
while os.path.exists(fname):
suffix_counter += 1
fname = original_fname.replace(
".dreem_metrics.h5", f"_{suffix_counter}.dreem_metrics.h5"
)
if suffix_counter > 0:
logger.info(f"File already exists. Saving to {fname} instead")
with h5py.File(fname, "a") as results_file:
# Create a group for this video
vid_group = results_file.require_group(vid_name)
# Save each metric
for metric_name, value in results.items():
if metric_name == "motmetrics":
# For num_switches, save mot_summary and mot_events separately
mot_summary = value[0]
mot_events = value[1]
frame_switch_map = value[2]
mot_summary_group = vid_group.require_group("mot_summary")
# Loop through each row in mot_summary and save as attributes
for _, row in mot_summary.iterrows():
mot_summary_group.attrs[row.name] = row["acc"]
# save extra metadata for frames in which there is a switch
for frame_id, switch in frame_switch_map.items():
frame = preds[frame_id]
frame = frame.to("cpu")
if switch:
_ = frame.to_h5(
vid_group,
frame.get_gt_track_ids().cpu().numpy(),
save={
"crop": True,
"features": True,
"embeddings": True,
},
)
else:
_ = frame.to_h5(
vid_group, frame.get_gt_track_ids().cpu().numpy()
)
# save motevents log to csv
motevents_path = os.path.join(
self.test_results["save_path"], f"{vid_name}.motevents.csv"
)
logger.info(f"Saving motevents log to {motevents_path}")
mot_events.to_csv(motevents_path, index=False)
elif metric_name == "global_tracking_accuracy":
gta_by_gt_track = value
gta_group = vid_group.require_group("global_tracking_accuracy")
# save as a key value pair with gt track id: gta
for gt_track_id, gta in gta_by_gt_track.items():
gta_group.attrs[f"track_{gt_track_id}"] = gta
# save the tracking results to a slp/labelled masks file
if isinstance(self.trainer.test_dataloaders.dataset, CellTrackingDataset):
outpath = os.path.join(
self.test_results["save_path"],
f"{vid_name}.dreem_inference.{datetime.now().strftime('%m-%d-%Y-%H-%M-%S')}.tif",
)
pred_imgs = []
for frame in preds:
frame_masks = []
for instance in frame.instances:
# centroid = instance.centroid["centroid"] # Currently unused but available if needed
mask = instance.mask.cpu().numpy()
track_id = instance.pred_track_id.cpu().numpy().item()
mask = mask.astype(np.uint8)
mask[mask != 0] = track_id # label the mask with the track id
frame_masks.append(mask)
frame_mask = np.max(frame_masks, axis=0)
pred_imgs.append(frame_mask)
pred_imgs = np.stack(pred_imgs)
tifffile.imwrite(outpath, pred_imgs.astype(np.uint16))
else:
outpath = os.path.join(
self.test_results["save_path"],
f"{vid_name}.dreem_inference.{datetime.now().strftime('%m-%d-%Y-%H-%M-%S')}.slp",
)
pred_slp = []
logger.info(f"Saving inference results to {outpath}")
# save the tracking results to a slp file
tracks = {}
for frame in preds:
if frame.frame_id.item() == 0:
video = (
sio.Video(frame.video)
if isinstance(frame.video, str)
else sio.Video
)
lf, tracks = frame.to_slp(tracks, video=video)
pred_slp.append(lf)
pred_slp = sio.Labels(pred_slp)
pred_slp.save(outpath)
# clear the preds
self.test_results["preds"] = []
__init__(model_cfg=None, tracker_cfg=None, loss_cfg=None, optimizer_cfg=None, scheduler_cfg=None, metrics=None, persistent_tracking=None, test_save_path='./test_results.h5')
¶
Initialize a lightning module for GTR.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
model_cfg
|
dict | None
|
hyperparameters for GlobalTrackingTransformer |
None
|
tracker_cfg
|
dict | None
|
The parameters used for the tracker post-processing |
None
|
loss_cfg
|
dict | None
|
hyperparameters for AssoLoss |
None
|
optimizer_cfg
|
dict | None
|
hyper parameters used for optimizer.
Only used to overwrite |
None
|
scheduler_cfg
|
dict | None
|
hyperparameters for lr_scheduler used to overwrite `configure_optimizer |
None
|
metrics
|
dict[str, list[str]] | None
|
a dict containing the metrics to be computed during train, val, and test. |
None
|
persistent_tracking
|
dict[str, bool] | None
|
a dict containing whether to use persistent tracking during train, val and test inference. |
None
|
test_save_path
|
str
|
path to a directory to save the eval and tracking results to |
'./test_results.h5'
|
Source code in dreem/models/gtr_runner.py
def __init__(
self,
model_cfg: dict | None = None,
tracker_cfg: dict | None = None,
loss_cfg: dict | None = None,
optimizer_cfg: dict | None = None,
scheduler_cfg: dict | None = None,
metrics: dict[str, list[str]] | None = None,
persistent_tracking: dict[str, bool] | None = None,
test_save_path: str = "./test_results.h5",
):
"""Initialize a lightning module for GTR.
Args:
model_cfg: hyperparameters for GlobalTrackingTransformer
tracker_cfg: The parameters used for the tracker post-processing
loss_cfg: hyperparameters for AssoLoss
optimizer_cfg: hyper parameters used for optimizer.
Only used to overwrite `configure_optimizer`
scheduler_cfg: hyperparameters for lr_scheduler used to overwrite `configure_optimizer
metrics: a dict containing the metrics to be computed during train, val, and test.
persistent_tracking: a dict containing whether to use persistent tracking during train, val and test inference.
test_save_path: path to a directory to save the eval and tracking results to
"""
super().__init__()
self.save_hyperparameters()
self.model_cfg = model_cfg if model_cfg else {}
self.loss_cfg = loss_cfg if loss_cfg else {}
self.tracker_cfg = tracker_cfg if tracker_cfg else {}
self.model = GlobalTrackingTransformer(**self.model_cfg)
self.loss = AssoLoss(**self.loss_cfg)
if self.tracker_cfg.get("tracker_type", "standard") == "batch":
from dreem.inference.batch_tracker import BatchTracker
self.tracker = BatchTracker(**self.tracker_cfg)
else:
from dreem.inference.tracker import Tracker
self.tracker = Tracker(**self.tracker_cfg)
self.optimizer_cfg = optimizer_cfg
self.scheduler_cfg = scheduler_cfg
self.metrics = metrics if metrics is not None else self.DEFAULT_METRICS
self.persistent_tracking = (
persistent_tracking
if persistent_tracking is not None
else self.DEFAULT_TRACKING
)
self.test_results = {"preds": [], "save_path": test_save_path}
configure_optimizers()
¶
Get optimizers and schedulers for training.
Is overridden by config but defaults to Adam + ReduceLROnPlateau.
Returns:
| Type | Description |
|---|---|
dict
|
an optimizer config dict containing the optimizer, scheduler, and scheduler params |
Source code in dreem/models/gtr_runner.py
def configure_optimizers(self) -> dict:
"""Get optimizers and schedulers for training.
Is overridden by config but defaults to Adam + ReduceLROnPlateau.
Returns:
an optimizer config dict containing the optimizer, scheduler, and scheduler params
"""
# todo: init from config
if self.optimizer_cfg is None:
optimizer = torch.optim.Adam(self.parameters(), lr=1e-4, betas=(0.9, 0.999))
else:
optimizer = init_optimizer(self.parameters(), self.optimizer_cfg)
if self.scheduler_cfg is None:
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
optimizer, "min", 0.5, 10
)
else:
scheduler = init_scheduler(optimizer, self.scheduler_cfg)
return {
"optimizer": optimizer,
"lr_scheduler": {
"scheduler": scheduler,
"monitor": "val_loss",
"interval": "epoch",
"frequency": 1,
},
}
forward(ref_instances, query_instances=None)
¶
Execute forward pass of the lightning module.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
ref_instances
|
list[Instance]
|
a list of |
required |
query_instances
|
list[Instance] | None
|
a list of |
None
|
Returns:
| Type | Description |
|---|---|
list[AssociationMatrix]
|
An association matrix between objects |
Source code in dreem/models/gtr_runner.py
def forward(
self,
ref_instances: list["Instance"],
query_instances: list["Instance"] | None = None,
) -> list["AssociationMatrix"]:
"""Execute forward pass of the lightning module.
Args:
ref_instances: a list of `Instance` objects containing crops and other data needed for transformer model
query_instances: a list of `Instance` objects used as queries in the decoder. Mostly used for inference.
Returns:
An association matrix between objects
"""
asso_preds = self.model(ref_instances, query_instances)
return asso_preds
log_metrics(result, batch_size, mode)
¶
Log metrics computed during evaluation.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
result
|
dict
|
A dict containing metrics to be logged. |
required |
batch_size
|
int
|
the size of the batch used to compute the metrics |
required |
mode
|
str
|
One of {'train', 'test' or 'val'}. Used as prefix while logging. |
required |
Source code in dreem/models/gtr_runner.py
def log_metrics(self, result: dict, batch_size: int, mode: str) -> None:
"""Log metrics computed during evaluation.
Args:
result: A dict containing metrics to be logged.
batch_size: the size of the batch used to compute the metrics
mode: One of {'train', 'test' or 'val'}. Used as prefix while logging.
"""
if result:
batch_size = result.pop("batch_size")
for metric, val in result.items():
if isinstance(val, torch.Tensor):
val = val.item()
self.log(f"{mode}_{metric}", val, batch_size=batch_size)
on_test_end()
¶
Run inference and metrics pipeline to compute metrics for test set.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
test_results
|
dict containing predictions and metrics to be filled out in metrics.evaluate |
required | |
metrics
|
list of metrics to compute |
required |
Source code in dreem/models/gtr_runner.py
def on_test_end(self):
"""Run inference and metrics pipeline to compute metrics for test set.
Args:
test_results: dict containing predictions and metrics to be filled out in metrics.evaluate
metrics: list of metrics to compute
"""
# input validation
metrics_to_compute = self.metrics[
"test"
] # list of metrics to compute, or "all"
if metrics_to_compute == "all":
metrics_to_compute = ["motmetrics", "global_tracking_accuracy"]
if isinstance(metrics_to_compute, str):
metrics_to_compute = [metrics_to_compute]
for metric in metrics_to_compute:
if metric not in ["motmetrics", "global_tracking_accuracy"]:
raise ValueError(
f"Metric {metric} not supported. Please select from 'motmetrics' or 'global_tracking_accuracy'"
)
preds = self.test_results["preds"]
# results is a dict with key being the metric name, and value being the metric value computed
results = metrics.evaluate(preds, metrics_to_compute)
# save metrics and frame metadata to hdf5
# Get the video name from the first frame
vid_name = Path(preds[0].vid_name).stem
# save the results to an hdf5 file
fname = os.path.join(
self.test_results["save_path"], f"{vid_name}.dreem_metrics.h5"
)
logger.info(f"Saving metrics to {fname}")
# Check if the h5 file exists and add a suffix to prevent name collision
suffix_counter = 0
original_fname = fname
while os.path.exists(fname):
suffix_counter += 1
fname = original_fname.replace(
".dreem_metrics.h5", f"_{suffix_counter}.dreem_metrics.h5"
)
if suffix_counter > 0:
logger.info(f"File already exists. Saving to {fname} instead")
with h5py.File(fname, "a") as results_file:
# Create a group for this video
vid_group = results_file.require_group(vid_name)
# Save each metric
for metric_name, value in results.items():
if metric_name == "motmetrics":
# For num_switches, save mot_summary and mot_events separately
mot_summary = value[0]
mot_events = value[1]
frame_switch_map = value[2]
mot_summary_group = vid_group.require_group("mot_summary")
# Loop through each row in mot_summary and save as attributes
for _, row in mot_summary.iterrows():
mot_summary_group.attrs[row.name] = row["acc"]
# save extra metadata for frames in which there is a switch
for frame_id, switch in frame_switch_map.items():
frame = preds[frame_id]
frame = frame.to("cpu")
if switch:
_ = frame.to_h5(
vid_group,
frame.get_gt_track_ids().cpu().numpy(),
save={
"crop": True,
"features": True,
"embeddings": True,
},
)
else:
_ = frame.to_h5(
vid_group, frame.get_gt_track_ids().cpu().numpy()
)
# save motevents log to csv
motevents_path = os.path.join(
self.test_results["save_path"], f"{vid_name}.motevents.csv"
)
logger.info(f"Saving motevents log to {motevents_path}")
mot_events.to_csv(motevents_path, index=False)
elif metric_name == "global_tracking_accuracy":
gta_by_gt_track = value
gta_group = vid_group.require_group("global_tracking_accuracy")
# save as a key value pair with gt track id: gta
for gt_track_id, gta in gta_by_gt_track.items():
gta_group.attrs[f"track_{gt_track_id}"] = gta
# save the tracking results to a slp/labelled masks file
if isinstance(self.trainer.test_dataloaders.dataset, CellTrackingDataset):
outpath = os.path.join(
self.test_results["save_path"],
f"{vid_name}.dreem_inference.{datetime.now().strftime('%m-%d-%Y-%H-%M-%S')}.tif",
)
pred_imgs = []
for frame in preds:
frame_masks = []
for instance in frame.instances:
# centroid = instance.centroid["centroid"] # Currently unused but available if needed
mask = instance.mask.cpu().numpy()
track_id = instance.pred_track_id.cpu().numpy().item()
mask = mask.astype(np.uint8)
mask[mask != 0] = track_id # label the mask with the track id
frame_masks.append(mask)
frame_mask = np.max(frame_masks, axis=0)
pred_imgs.append(frame_mask)
pred_imgs = np.stack(pred_imgs)
tifffile.imwrite(outpath, pred_imgs.astype(np.uint16))
else:
outpath = os.path.join(
self.test_results["save_path"],
f"{vid_name}.dreem_inference.{datetime.now().strftime('%m-%d-%Y-%H-%M-%S')}.slp",
)
pred_slp = []
logger.info(f"Saving inference results to {outpath}")
# save the tracking results to a slp file
tracks = {}
for frame in preds:
if frame.frame_id.item() == 0:
video = (
sio.Video(frame.video)
if isinstance(frame.video, str)
else sio.Video
)
lf, tracks = frame.to_slp(tracks, video=video)
pred_slp.append(lf)
pred_slp = sio.Labels(pred_slp)
pred_slp.save(outpath)
# clear the preds
self.test_results["preds"] = []
on_validation_epoch_end()
¶
Execute hook for validation end.
Currently, we simply clear the gpu cache and do garbage collection.
predict_step(batch, batch_idx)
¶
Run inference for model.
Computes association + assignment.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
batch
|
list[list[Frame]]
|
A single batch from the dataset which is a list of |
required |
batch_idx
|
int
|
the batch number used by lightning |
required |
Returns:
| Type | Description |
|---|---|
list[Frame]
|
A list of dicts where each dict is a frame containing the predicted track ids |
Source code in dreem/models/gtr_runner.py
def predict_step(self, batch: list[list["Frame"]], batch_idx: int) -> list["Frame"]:
"""Run inference for model.
Computes association + assignment.
Args:
batch: A single batch from the dataset which is a list of `Frame` objects
with length `clip_length` containing Instances and other metadata.
batch_idx: the batch number used by lightning
Returns:
A list of dicts where each dict is a frame containing the predicted track ids
"""
frames_pred = self.tracker(self.model, batch[0])
return frames_pred
test_step(test_batch, batch_idx)
¶
Execute single test step for model.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
test_batch
|
list[list[Frame]]
|
A single batch from the dataset which is a list of |
required |
batch_idx
|
int
|
the batch number used by lightning |
required |
Returns:
| Type | Description |
|---|---|
dict[str, float]
|
A dict containing the val loss plus any other metrics specified |
Source code in dreem/models/gtr_runner.py
def test_step(
self, test_batch: list[list["Frame"]], batch_idx: int
) -> dict[str, float]:
"""Execute single test step for model.
Args:
test_batch: A single batch from the dataset which is a list of `Frame` objects
with length `clip_length` containing Instances and other metadata.
batch_idx: the batch number used by lightning
Returns:
A dict containing the val loss plus any other metrics specified
"""
result = self._shared_eval_step(test_batch[0], mode="test")
self.log_metrics(result, len(test_batch[0]), "test")
return result
training_step(train_batch, batch_idx)
¶
Execute single training step for model.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
train_batch
|
list[list[Frame]]
|
A single batch from the dataset which is a list of |
required |
batch_idx
|
int
|
the batch number used by lightning |
required |
Returns:
| Type | Description |
|---|---|
dict[str, float]
|
A dict containing the train loss plus any other metrics specified |
Source code in dreem/models/gtr_runner.py
def training_step(
self, train_batch: list[list["Frame"]], batch_idx: int
) -> dict[str, float]:
"""Execute single training step for model.
Args:
train_batch: A single batch from the dataset which is a list of `Frame` objects
with length `clip_length` containing Instances and other metadata.
batch_idx: the batch number used by lightning
Returns:
A dict containing the train loss plus any other metrics specified
"""
result = self._shared_eval_step(train_batch[0], mode="train")
self.log_metrics(result, len(train_batch[0]), "train")
return result
validation_step(val_batch, batch_idx)
¶
Execute single val step for model.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
val_batch
|
list[list[Frame]]
|
A single batch from the dataset which is a list of |
required |
batch_idx
|
int
|
the batch number used by lightning |
required |
Returns:
| Type | Description |
|---|---|
dict[str, float]
|
A dict containing the val loss plus any other metrics specified |
Source code in dreem/models/gtr_runner.py
def validation_step(
self, val_batch: list[list["Frame"]], batch_idx: int
) -> dict[str, float]:
"""Execute single val step for model.
Args:
val_batch: A single batch from the dataset which is a list of `Frame` objects
with length `clip_length` containing Instances and other metadata.
batch_idx: the batch number used by lightning
Returns:
A dict containing the val loss plus any other metrics specified
"""
result = self._shared_eval_step(val_batch[0], mode="val")
self.log_metrics(result, len(val_batch[0]), "val")
return result
GlobalTrackingTransformer
¶
Bases: Module
Modular GTR model composed of visual encoder + transformer used for tracking.
Methods:
| Name | Description |
|---|---|
__init__ |
Initialize GTR. |
extract_features |
Extract features from instances using visual encoder backbone. |
forward |
Execute forward pass of GTR Model to get asso matrix. |
Source code in dreem/models/global_tracking_transformer.py
class GlobalTrackingTransformer(torch.nn.Module):
"""Modular GTR model composed of visual encoder + transformer used for tracking."""
def __init__(
self,
encoder_cfg: dict | None = None,
d_model: int = 1024,
nhead: int = 8,
num_encoder_layers: int = 6,
num_decoder_layers: int = 6,
dropout: int = 0.1,
activation: str = "relu",
return_intermediate_dec: bool = False,
norm: bool = False,
num_layers_attn_head: int = 2,
dropout_attn_head: int = 0.1,
embedding_meta: dict | None = None,
return_embedding: bool = False,
decoder_self_attn: bool = False,
):
"""Initialize GTR.
Args:
encoder_cfg: Dictionary of arguments to pass to the CNN constructor,
e.g: `cfg = {"model_name": "resnet18", "pretrained": False, "in_chans": 3}`
d_model: The number of features in the encoder/decoder inputs.
nhead: The number of heads in the transformer encoder/decoder.
num_encoder_layers: The number of encoder-layers in the encoder.
num_decoder_layers: The number of decoder-layers in the decoder.
dropout: Dropout value applied to the output of transformer layers.
activation: Activation function to use.
return_intermediate_dec: Return intermediate layers from decoder.
norm: If True, normalize output of encoder and decoder.
num_layers_attn_head: The number of layers in the attention head.
dropout_attn_head: Dropout value for the attention_head.
embedding_meta: Metadata for positional embeddings. See below.
return_embedding: Whether to return the positional embeddings
decoder_self_attn: If True, use decoder self attention.
More details on `embedding_meta`:
By default this will be an empty dict and indicate
that no positional embeddings should be used. To use the positional embeddings
pass in a dictionary containing a "pos" and "temp" key with subdictionaries for correct parameters ie:
`{"pos": {'mode': 'learned', 'emb_num': 16, 'over_boxes: True},
"temp": {'mode': 'learned', 'emb_num': 16}}`. (see `dreem.models.embeddings.Embedding.EMB_TYPES`
and `dreem.models.embeddings.Embedding.EMB_MODES` for embedding parameters).
"""
super().__init__()
if not encoder_cfg:
encoder_cfg = {}
self.visual_encoder = create_visual_encoder(d_model=d_model, **encoder_cfg)
self.transformer = Transformer(
d_model=d_model,
nhead=nhead,
num_encoder_layers=num_encoder_layers,
num_decoder_layers=num_decoder_layers,
dropout=dropout,
activation=activation,
return_intermediate_dec=return_intermediate_dec,
norm=norm,
num_layers_attn_head=num_layers_attn_head,
dropout_attn_head=dropout_attn_head,
embedding_meta=embedding_meta,
return_embedding=return_embedding,
decoder_self_attn=decoder_self_attn,
encoder_cfg=encoder_cfg,
)
def forward(
self, ref_instances: list["Instance"], query_instances: list["Instance"] = None
) -> list["AssociationMatrix"]:
"""Execute forward pass of GTR Model to get asso matrix.
Args:
ref_instances: List of instances from chunk containing crops of objects + gt label info
query_instances: list of instances used as query in decoder.
Returns:
An N_T x N association matrix
"""
# Extract feature representations with pre-trained encoder.
self.extract_features(ref_instances)
if query_instances:
self.extract_features(query_instances)
asso_preds = self.transformer(ref_instances, query_instances)
return asso_preds
def extract_features(
self, instances: list["Instance"], force_recompute: bool = False
) -> None:
"""Extract features from instances using visual encoder backbone.
Args:
instances: A list of instances to compute features for
force_recompute: indicate whether to compute features for all instances regardless of if they have instances
"""
if not force_recompute:
instances_to_compute = [
instance
for instance in instances
if instance.has_crop() and not instance.has_features()
]
else:
instances_to_compute = instances
if len(instances_to_compute) == 0:
return
elif len(instances_to_compute) == 1: # handle batch norm error when B=1
instances_to_compute = instances
crops = torch.concatenate([instance.crop for instance in instances_to_compute])
features = self.visual_encoder(crops)
features = features.to(device=instances_to_compute[0].device)
for i, z_i in enumerate(features):
instances_to_compute[i].features = z_i
__init__(encoder_cfg=None, d_model=1024, nhead=8, num_encoder_layers=6, num_decoder_layers=6, dropout=0.1, activation='relu', return_intermediate_dec=False, norm=False, num_layers_attn_head=2, dropout_attn_head=0.1, embedding_meta=None, return_embedding=False, decoder_self_attn=False)
¶
Initialize GTR.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
encoder_cfg
|
dict | None
|
Dictionary of arguments to pass to the CNN constructor,
e.g: |
None
|
d_model
|
int
|
The number of features in the encoder/decoder inputs. |
1024
|
nhead
|
int
|
The number of heads in the transformer encoder/decoder. |
8
|
num_encoder_layers
|
int
|
The number of encoder-layers in the encoder. |
6
|
num_decoder_layers
|
int
|
The number of decoder-layers in the decoder. |
6
|
dropout
|
int
|
Dropout value applied to the output of transformer layers. |
0.1
|
activation
|
str
|
Activation function to use. |
'relu'
|
return_intermediate_dec
|
bool
|
Return intermediate layers from decoder. |
False
|
norm
|
bool
|
If True, normalize output of encoder and decoder. |
False
|
num_layers_attn_head
|
int
|
The number of layers in the attention head. |
2
|
dropout_attn_head
|
int
|
Dropout value for the attention_head. |
0.1
|
embedding_meta
|
dict | None
|
Metadata for positional embeddings. See below. |
None
|
return_embedding
|
bool
|
Whether to return the positional embeddings |
False
|
decoder_self_attn
|
bool
|
If True, use decoder self attention. More details on |
False
|
Source code in dreem/models/global_tracking_transformer.py
def __init__(
self,
encoder_cfg: dict | None = None,
d_model: int = 1024,
nhead: int = 8,
num_encoder_layers: int = 6,
num_decoder_layers: int = 6,
dropout: int = 0.1,
activation: str = "relu",
return_intermediate_dec: bool = False,
norm: bool = False,
num_layers_attn_head: int = 2,
dropout_attn_head: int = 0.1,
embedding_meta: dict | None = None,
return_embedding: bool = False,
decoder_self_attn: bool = False,
):
"""Initialize GTR.
Args:
encoder_cfg: Dictionary of arguments to pass to the CNN constructor,
e.g: `cfg = {"model_name": "resnet18", "pretrained": False, "in_chans": 3}`
d_model: The number of features in the encoder/decoder inputs.
nhead: The number of heads in the transformer encoder/decoder.
num_encoder_layers: The number of encoder-layers in the encoder.
num_decoder_layers: The number of decoder-layers in the decoder.
dropout: Dropout value applied to the output of transformer layers.
activation: Activation function to use.
return_intermediate_dec: Return intermediate layers from decoder.
norm: If True, normalize output of encoder and decoder.
num_layers_attn_head: The number of layers in the attention head.
dropout_attn_head: Dropout value for the attention_head.
embedding_meta: Metadata for positional embeddings. See below.
return_embedding: Whether to return the positional embeddings
decoder_self_attn: If True, use decoder self attention.
More details on `embedding_meta`:
By default this will be an empty dict and indicate
that no positional embeddings should be used. To use the positional embeddings
pass in a dictionary containing a "pos" and "temp" key with subdictionaries for correct parameters ie:
`{"pos": {'mode': 'learned', 'emb_num': 16, 'over_boxes: True},
"temp": {'mode': 'learned', 'emb_num': 16}}`. (see `dreem.models.embeddings.Embedding.EMB_TYPES`
and `dreem.models.embeddings.Embedding.EMB_MODES` for embedding parameters).
"""
super().__init__()
if not encoder_cfg:
encoder_cfg = {}
self.visual_encoder = create_visual_encoder(d_model=d_model, **encoder_cfg)
self.transformer = Transformer(
d_model=d_model,
nhead=nhead,
num_encoder_layers=num_encoder_layers,
num_decoder_layers=num_decoder_layers,
dropout=dropout,
activation=activation,
return_intermediate_dec=return_intermediate_dec,
norm=norm,
num_layers_attn_head=num_layers_attn_head,
dropout_attn_head=dropout_attn_head,
embedding_meta=embedding_meta,
return_embedding=return_embedding,
decoder_self_attn=decoder_self_attn,
encoder_cfg=encoder_cfg,
)
extract_features(instances, force_recompute=False)
¶
Extract features from instances using visual encoder backbone.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
instances
|
list[Instance]
|
A list of instances to compute features for |
required |
force_recompute
|
bool
|
indicate whether to compute features for all instances regardless of if they have instances |
False
|
Source code in dreem/models/global_tracking_transformer.py
def extract_features(
self, instances: list["Instance"], force_recompute: bool = False
) -> None:
"""Extract features from instances using visual encoder backbone.
Args:
instances: A list of instances to compute features for
force_recompute: indicate whether to compute features for all instances regardless of if they have instances
"""
if not force_recompute:
instances_to_compute = [
instance
for instance in instances
if instance.has_crop() and not instance.has_features()
]
else:
instances_to_compute = instances
if len(instances_to_compute) == 0:
return
elif len(instances_to_compute) == 1: # handle batch norm error when B=1
instances_to_compute = instances
crops = torch.concatenate([instance.crop for instance in instances_to_compute])
features = self.visual_encoder(crops)
features = features.to(device=instances_to_compute[0].device)
for i, z_i in enumerate(features):
instances_to_compute[i].features = z_i
forward(ref_instances, query_instances=None)
¶
Execute forward pass of GTR Model to get asso matrix.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
ref_instances
|
list[Instance]
|
List of instances from chunk containing crops of objects + gt label info |
required |
query_instances
|
list[Instance]
|
list of instances used as query in decoder. |
None
|
Returns:
| Type | Description |
|---|---|
list[AssociationMatrix]
|
An N_T x N association matrix |
Source code in dreem/models/global_tracking_transformer.py
def forward(
self, ref_instances: list["Instance"], query_instances: list["Instance"] = None
) -> list["AssociationMatrix"]:
"""Execute forward pass of GTR Model to get asso matrix.
Args:
ref_instances: List of instances from chunk containing crops of objects + gt label info
query_instances: list of instances used as query in decoder.
Returns:
An N_T x N association matrix
"""
# Extract feature representations with pre-trained encoder.
self.extract_features(ref_instances)
if query_instances:
self.extract_features(query_instances)
asso_preds = self.transformer(ref_instances, query_instances)
return asso_preds
Transformer
¶
Bases: Module
Transformer class.
Methods:
| Name | Description |
|---|---|
__init__ |
Initialize Transformer. |
forward |
Execute a forward pass through the transformer and attention head. |
Source code in dreem/models/transformer.py
class Transformer(torch.nn.Module):
"""Transformer class."""
def __init__(
self,
d_model: int = 1024,
nhead: int = 8,
num_encoder_layers: int = 6,
num_decoder_layers: int = 6,
dropout: float = 0.1,
activation: str = "relu",
return_intermediate_dec: bool = False,
norm: bool = False,
num_layers_attn_head: int = 2,
dropout_attn_head: float = 0.1,
embedding_meta: dict | None = None,
return_embedding: bool = False,
decoder_self_attn: bool = False,
encoder_cfg: dict | None = None,
) -> None:
"""Initialize Transformer.
Args:
d_model: The number of features in the encoder/decoder inputs.
nhead: The number of heads in the transformer encoder/decoder.
num_encoder_layers: The number of encoder-layers in the encoder.
num_decoder_layers: The number of decoder-layers in the decoder.
dropout: Dropout value applied to the output of transformer layers.
activation: Activation function to use.
return_intermediate_dec: Return intermediate layers from decoder.
norm: If True, normalize output of encoder and decoder.
num_layers_attn_head: The number of layers in the attention head.
dropout_attn_head: Dropout value for the attention_head.
embedding_meta: Metadata for positional embeddings. See below.
return_embedding: Whether to return the positional embeddings
decoder_self_attn: If True, use decoder self attention.
encoder_cfg: Encoder configuration.
More details on `embedding_meta`:
By default this will be an empty dict and indicate
that no positional embeddings should be used. To use the positional embeddings
pass in a dictionary containing a "pos" and "temp" key with subdictionaries for correct parameters ie:
{"pos": {'mode': 'learned', 'emb_num': 16, 'over_boxes: 'True'},
"temp": {'mode': 'learned', 'emb_num': 16}}. (see `dreem.models.embeddings.Embedding.EMB_TYPES`
and `dreem.models.embeddings.Embedding.EMB_MODES` for embedding parameters).
"""
super().__init__()
self.d_model = dim_feedforward = feature_dim_attn_head = d_model
self.embedding_meta = embedding_meta
self.return_embedding = return_embedding
self.encoder_cfg = encoder_cfg
self.pos_emb = Embedding(emb_type="off", mode="off", features=self.d_model)
self.temp_emb = Embedding(emb_type="off", mode="off", features=self.d_model)
if self.embedding_meta:
if "pos" in self.embedding_meta:
pos_emb_cfg = self.embedding_meta["pos"]
if pos_emb_cfg:
self.pos_emb = Embedding(
emb_type="pos", features=self.d_model, **pos_emb_cfg
)
if "temp" in self.embedding_meta:
temp_emb_cfg = self.embedding_meta["temp"]
if temp_emb_cfg:
self.temp_emb = Embedding(
emb_type="temp", features=self.d_model, **temp_emb_cfg
)
self.fourier_embeddings = FourierPositionalEmbeddings(
n_components=8, d_model=d_model
)
# Transformer Encoder
encoder_layer = TransformerEncoderLayer(
d_model, nhead, dim_feedforward, dropout, activation, norm
)
encoder_norm = nn.LayerNorm(d_model) if (norm) else None
# only used if using descriptor visual encoder; default resnet encoder uses d_model directly
if self.encoder_cfg and "encoder_type" in self.encoder_cfg:
self.visual_feat_dim = (
self.encoder_cfg["ndim"] if "ndim" in self.encoder_cfg else 5
) # 5 is default for descriptor
self.fourier_proj = nn.Linear(self.d_model + self.visual_feat_dim, d_model)
self.fourier_norm = nn.LayerNorm(self.d_model)
self.encoder = TransformerEncoder(
encoder_layer, num_encoder_layers, encoder_norm
)
# Transformer Decoder
decoder_layer = TransformerDecoderLayer(
d_model,
nhead,
dim_feedforward,
dropout,
activation,
norm,
decoder_self_attn,
)
decoder_norm = nn.LayerNorm(d_model) if (norm) else None
self.decoder = TransformerDecoder(
decoder_layer, num_decoder_layers, return_intermediate_dec, decoder_norm
)
# Transformer attention head
self.attn_head = ATTWeightHead(
feature_dim=feature_dim_attn_head,
num_layers=num_layers_attn_head,
dropout=dropout_attn_head,
)
self._reset_parameters()
def _reset_parameters(self):
"""Initialize model weights from xavier distribution."""
for p in self.parameters():
if not torch.nn.parameter.is_lazy(p) and p.dim() > 1:
try:
nn.init.xavier_uniform_(p)
except ValueError as e:
print(f"Failed Trying to initialize {p}")
raise (e)
def forward(
self,
ref_instances: list[Instance],
query_instances: list[Instance] | None = None,
) -> list[AssociationMatrix]:
"""Execute a forward pass through the transformer and attention head.
Args:
ref_instances: A list of instance objects (See `dreem.io.Instance` for more info.)
query_instances: An set of instances to be used as decoder queries.
Returns:
asso_output: A list of torch.Tensors of shape (L, n_query, total_instances) where:
L: number of decoder blocks
n_query: number of instances in current query/frame
total_instances: number of instances in window
"""
ref_features = torch.cat(
[instance.features for instance in ref_instances], dim=0
).unsqueeze(0)
# window_length = len(frames)
# instances_per_frame = [frame.num_detected for frame in frames]
total_instances = len(ref_instances)
embed_dim = self.d_model
# print(f'T: {window_length}; N: {total_instances}; N_t: {instances_per_frame} n_reid: {reid_features.shape}')
ref_boxes = get_boxes(ref_instances) # total_instances, 4
ref_boxes = torch.nan_to_num(ref_boxes, -1.0)
ref_times, query_times = get_times(ref_instances, query_instances)
# window_length = len(ref_times.unique()) # Currently unused but may be useful for debugging
ref_temp_emb = self.temp_emb(ref_times)
ref_pos_emb = self.pos_emb(ref_boxes)
if self.return_embedding:
for i, instance in enumerate(ref_instances):
instance.add_embedding("pos", ref_pos_emb[i])
instance.add_embedding("temp", ref_temp_emb[i])
ref_emb = (ref_pos_emb + ref_temp_emb) / 2.0
ref_emb = ref_emb.view(1, total_instances, embed_dim)
ref_emb = ref_emb.permute(1, 0, 2) # (total_instances, batch_size, embed_dim)
batch_size, total_instances = ref_features.shape[:-1]
ref_features = ref_features.permute(
1, 0, 2
) # (total_instances, batch_size, embed_dim)
encoder_queries = ref_features
# apply fourier embeddings if using fourier rope, OR if using descriptor (compact) visual encoder
if (
self.embedding_meta
and "use_fourier" in self.embedding_meta
and self.embedding_meta["use_fourier"]
) or (
self.encoder_cfg
and "encoder_type" in self.encoder_cfg
and self.encoder_cfg["encoder_type"] == "descriptor"
):
encoder_queries = apply_fourier_embeddings(
encoder_queries,
ref_times,
self.d_model,
self.fourier_embeddings,
self.fourier_proj,
self.fourier_norm,
)
encoder_features = self.encoder(
encoder_queries, pos_emb=ref_emb
) # (total_instances, batch_size, embed_dim)
n_query = total_instances
query_features = ref_features
query_pos_emb = ref_pos_emb
query_temp_emb = ref_temp_emb
query_emb = ref_emb
if query_instances is not None:
n_query = len(query_instances)
query_features = torch.cat(
[instance.features for instance in query_instances], dim=0
).unsqueeze(0)
query_features = query_features.permute(
1, 0, 2
) # (n_query, batch_size, embed_dim)
query_boxes = get_boxes(query_instances)
query_boxes = torch.nan_to_num(query_boxes, -1.0)
query_temp_emb = self.temp_emb(query_times)
query_pos_emb = self.pos_emb(query_boxes)
query_emb = (query_pos_emb + query_temp_emb) / 2.0
query_emb = query_emb.view(1, n_query, embed_dim)
query_emb = query_emb.permute(1, 0, 2) # (n_query, batch_size, embed_dim)
else:
query_instances = ref_instances
query_times = ref_times
if self.return_embedding:
for i, instance in enumerate(query_instances):
instance.add_embedding("pos", query_pos_emb[i])
instance.add_embedding("temp", query_temp_emb[i])
# apply fourier embeddings if using fourier rope, OR if using descriptor (compact) visual encoder
if (
self.embedding_meta
and "use_fourier" in self.embedding_meta
and self.embedding_meta["use_fourier"]
) or (
self.encoder_cfg
and "encoder_type" in self.encoder_cfg
and self.encoder_cfg["encoder_type"] == "descriptor"
):
query_features = apply_fourier_embeddings(
query_features,
query_times,
self.d_model,
self.fourier_embeddings,
self.fourier_proj,
self.fourier_norm,
)
decoder_features = self.decoder(
query_features,
encoder_features,
ref_pos_emb=ref_emb,
query_pos_emb=query_emb,
) # (L, n_query, batch_size, embed_dim)
decoder_features = decoder_features.transpose(
1, 2
) # # (L, batch_size, n_query, embed_dim)
encoder_features = encoder_features.permute(1, 0, 2).view(
batch_size, total_instances, embed_dim
) # (batch_size, total_instances, embed_dim)
asso_output = []
for frame_features in decoder_features:
asso_matrix = self.attn_head(frame_features, encoder_features).view(
n_query, total_instances
)
asso_matrix = AssociationMatrix(asso_matrix, ref_instances, query_instances)
asso_output.append(asso_matrix)
# (L=1, n_query, total_instances)
return asso_output
__init__(d_model=1024, nhead=8, num_encoder_layers=6, num_decoder_layers=6, dropout=0.1, activation='relu', return_intermediate_dec=False, norm=False, num_layers_attn_head=2, dropout_attn_head=0.1, embedding_meta=None, return_embedding=False, decoder_self_attn=False, encoder_cfg=None)
¶
Initialize Transformer.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
d_model
|
int
|
The number of features in the encoder/decoder inputs. |
1024
|
nhead
|
int
|
The number of heads in the transformer encoder/decoder. |
8
|
num_encoder_layers
|
int
|
The number of encoder-layers in the encoder. |
6
|
num_decoder_layers
|
int
|
The number of decoder-layers in the decoder. |
6
|
dropout
|
float
|
Dropout value applied to the output of transformer layers. |
0.1
|
activation
|
str
|
Activation function to use. |
'relu'
|
return_intermediate_dec
|
bool
|
Return intermediate layers from decoder. |
False
|
norm
|
bool
|
If True, normalize output of encoder and decoder. |
False
|
num_layers_attn_head
|
int
|
The number of layers in the attention head. |
2
|
dropout_attn_head
|
float
|
Dropout value for the attention_head. |
0.1
|
embedding_meta
|
dict | None
|
Metadata for positional embeddings. See below. |
None
|
return_embedding
|
bool
|
Whether to return the positional embeddings |
False
|
decoder_self_attn
|
bool
|
If True, use decoder self attention. |
False
|
encoder_cfg
|
dict | None
|
Encoder configuration. More details on |
None
|
Source code in dreem/models/transformer.py
def __init__(
self,
d_model: int = 1024,
nhead: int = 8,
num_encoder_layers: int = 6,
num_decoder_layers: int = 6,
dropout: float = 0.1,
activation: str = "relu",
return_intermediate_dec: bool = False,
norm: bool = False,
num_layers_attn_head: int = 2,
dropout_attn_head: float = 0.1,
embedding_meta: dict | None = None,
return_embedding: bool = False,
decoder_self_attn: bool = False,
encoder_cfg: dict | None = None,
) -> None:
"""Initialize Transformer.
Args:
d_model: The number of features in the encoder/decoder inputs.
nhead: The number of heads in the transformer encoder/decoder.
num_encoder_layers: The number of encoder-layers in the encoder.
num_decoder_layers: The number of decoder-layers in the decoder.
dropout: Dropout value applied to the output of transformer layers.
activation: Activation function to use.
return_intermediate_dec: Return intermediate layers from decoder.
norm: If True, normalize output of encoder and decoder.
num_layers_attn_head: The number of layers in the attention head.
dropout_attn_head: Dropout value for the attention_head.
embedding_meta: Metadata for positional embeddings. See below.
return_embedding: Whether to return the positional embeddings
decoder_self_attn: If True, use decoder self attention.
encoder_cfg: Encoder configuration.
More details on `embedding_meta`:
By default this will be an empty dict and indicate
that no positional embeddings should be used. To use the positional embeddings
pass in a dictionary containing a "pos" and "temp" key with subdictionaries for correct parameters ie:
{"pos": {'mode': 'learned', 'emb_num': 16, 'over_boxes: 'True'},
"temp": {'mode': 'learned', 'emb_num': 16}}. (see `dreem.models.embeddings.Embedding.EMB_TYPES`
and `dreem.models.embeddings.Embedding.EMB_MODES` for embedding parameters).
"""
super().__init__()
self.d_model = dim_feedforward = feature_dim_attn_head = d_model
self.embedding_meta = embedding_meta
self.return_embedding = return_embedding
self.encoder_cfg = encoder_cfg
self.pos_emb = Embedding(emb_type="off", mode="off", features=self.d_model)
self.temp_emb = Embedding(emb_type="off", mode="off", features=self.d_model)
if self.embedding_meta:
if "pos" in self.embedding_meta:
pos_emb_cfg = self.embedding_meta["pos"]
if pos_emb_cfg:
self.pos_emb = Embedding(
emb_type="pos", features=self.d_model, **pos_emb_cfg
)
if "temp" in self.embedding_meta:
temp_emb_cfg = self.embedding_meta["temp"]
if temp_emb_cfg:
self.temp_emb = Embedding(
emb_type="temp", features=self.d_model, **temp_emb_cfg
)
self.fourier_embeddings = FourierPositionalEmbeddings(
n_components=8, d_model=d_model
)
# Transformer Encoder
encoder_layer = TransformerEncoderLayer(
d_model, nhead, dim_feedforward, dropout, activation, norm
)
encoder_norm = nn.LayerNorm(d_model) if (norm) else None
# only used if using descriptor visual encoder; default resnet encoder uses d_model directly
if self.encoder_cfg and "encoder_type" in self.encoder_cfg:
self.visual_feat_dim = (
self.encoder_cfg["ndim"] if "ndim" in self.encoder_cfg else 5
) # 5 is default for descriptor
self.fourier_proj = nn.Linear(self.d_model + self.visual_feat_dim, d_model)
self.fourier_norm = nn.LayerNorm(self.d_model)
self.encoder = TransformerEncoder(
encoder_layer, num_encoder_layers, encoder_norm
)
# Transformer Decoder
decoder_layer = TransformerDecoderLayer(
d_model,
nhead,
dim_feedforward,
dropout,
activation,
norm,
decoder_self_attn,
)
decoder_norm = nn.LayerNorm(d_model) if (norm) else None
self.decoder = TransformerDecoder(
decoder_layer, num_decoder_layers, return_intermediate_dec, decoder_norm
)
# Transformer attention head
self.attn_head = ATTWeightHead(
feature_dim=feature_dim_attn_head,
num_layers=num_layers_attn_head,
dropout=dropout_attn_head,
)
self._reset_parameters()
forward(ref_instances, query_instances=None)
¶
Execute a forward pass through the transformer and attention head.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
ref_instances
|
list[Instance]
|
A list of instance objects (See |
required |
query_instances
|
list[Instance] | None
|
An set of instances to be used as decoder queries. |
None
|
Returns:
| Name | Type | Description |
|---|---|---|
asso_output |
list[AssociationMatrix]
|
A list of torch.Tensors of shape (L, n_query, total_instances) where: L: number of decoder blocks n_query: number of instances in current query/frame total_instances: number of instances in window |
Source code in dreem/models/transformer.py
def forward(
self,
ref_instances: list[Instance],
query_instances: list[Instance] | None = None,
) -> list[AssociationMatrix]:
"""Execute a forward pass through the transformer and attention head.
Args:
ref_instances: A list of instance objects (See `dreem.io.Instance` for more info.)
query_instances: An set of instances to be used as decoder queries.
Returns:
asso_output: A list of torch.Tensors of shape (L, n_query, total_instances) where:
L: number of decoder blocks
n_query: number of instances in current query/frame
total_instances: number of instances in window
"""
ref_features = torch.cat(
[instance.features for instance in ref_instances], dim=0
).unsqueeze(0)
# window_length = len(frames)
# instances_per_frame = [frame.num_detected for frame in frames]
total_instances = len(ref_instances)
embed_dim = self.d_model
# print(f'T: {window_length}; N: {total_instances}; N_t: {instances_per_frame} n_reid: {reid_features.shape}')
ref_boxes = get_boxes(ref_instances) # total_instances, 4
ref_boxes = torch.nan_to_num(ref_boxes, -1.0)
ref_times, query_times = get_times(ref_instances, query_instances)
# window_length = len(ref_times.unique()) # Currently unused but may be useful for debugging
ref_temp_emb = self.temp_emb(ref_times)
ref_pos_emb = self.pos_emb(ref_boxes)
if self.return_embedding:
for i, instance in enumerate(ref_instances):
instance.add_embedding("pos", ref_pos_emb[i])
instance.add_embedding("temp", ref_temp_emb[i])
ref_emb = (ref_pos_emb + ref_temp_emb) / 2.0
ref_emb = ref_emb.view(1, total_instances, embed_dim)
ref_emb = ref_emb.permute(1, 0, 2) # (total_instances, batch_size, embed_dim)
batch_size, total_instances = ref_features.shape[:-1]
ref_features = ref_features.permute(
1, 0, 2
) # (total_instances, batch_size, embed_dim)
encoder_queries = ref_features
# apply fourier embeddings if using fourier rope, OR if using descriptor (compact) visual encoder
if (
self.embedding_meta
and "use_fourier" in self.embedding_meta
and self.embedding_meta["use_fourier"]
) or (
self.encoder_cfg
and "encoder_type" in self.encoder_cfg
and self.encoder_cfg["encoder_type"] == "descriptor"
):
encoder_queries = apply_fourier_embeddings(
encoder_queries,
ref_times,
self.d_model,
self.fourier_embeddings,
self.fourier_proj,
self.fourier_norm,
)
encoder_features = self.encoder(
encoder_queries, pos_emb=ref_emb
) # (total_instances, batch_size, embed_dim)
n_query = total_instances
query_features = ref_features
query_pos_emb = ref_pos_emb
query_temp_emb = ref_temp_emb
query_emb = ref_emb
if query_instances is not None:
n_query = len(query_instances)
query_features = torch.cat(
[instance.features for instance in query_instances], dim=0
).unsqueeze(0)
query_features = query_features.permute(
1, 0, 2
) # (n_query, batch_size, embed_dim)
query_boxes = get_boxes(query_instances)
query_boxes = torch.nan_to_num(query_boxes, -1.0)
query_temp_emb = self.temp_emb(query_times)
query_pos_emb = self.pos_emb(query_boxes)
query_emb = (query_pos_emb + query_temp_emb) / 2.0
query_emb = query_emb.view(1, n_query, embed_dim)
query_emb = query_emb.permute(1, 0, 2) # (n_query, batch_size, embed_dim)
else:
query_instances = ref_instances
query_times = ref_times
if self.return_embedding:
for i, instance in enumerate(query_instances):
instance.add_embedding("pos", query_pos_emb[i])
instance.add_embedding("temp", query_temp_emb[i])
# apply fourier embeddings if using fourier rope, OR if using descriptor (compact) visual encoder
if (
self.embedding_meta
and "use_fourier" in self.embedding_meta
and self.embedding_meta["use_fourier"]
) or (
self.encoder_cfg
and "encoder_type" in self.encoder_cfg
and self.encoder_cfg["encoder_type"] == "descriptor"
):
query_features = apply_fourier_embeddings(
query_features,
query_times,
self.d_model,
self.fourier_embeddings,
self.fourier_proj,
self.fourier_norm,
)
decoder_features = self.decoder(
query_features,
encoder_features,
ref_pos_emb=ref_emb,
query_pos_emb=query_emb,
) # (L, n_query, batch_size, embed_dim)
decoder_features = decoder_features.transpose(
1, 2
) # # (L, batch_size, n_query, embed_dim)
encoder_features = encoder_features.permute(1, 0, 2).view(
batch_size, total_instances, embed_dim
) # (batch_size, total_instances, embed_dim)
asso_output = []
for frame_features in decoder_features:
asso_matrix = self.attn_head(frame_features, encoder_features).view(
n_query, total_instances
)
asso_matrix = AssociationMatrix(asso_matrix, ref_instances, query_instances)
asso_output.append(asso_matrix)
# (L=1, n_query, total_instances)
return asso_output
VisualEncoder
¶
Bases: Module
Class wrapping around a visual feature extractor backbone.
Currently CNN only.
Methods:
| Name | Description |
|---|---|
__init__ |
Initialize Visual Encoder. |
encoder_dim |
Compute dummy forward pass of encoder model and get embedding dimension. |
forward |
Forward pass of feature extractor to get feature vector. |
select_feature_extractor |
Select the appropriate feature extractor based on config. |
Source code in dreem/models/visual_encoder.py
class VisualEncoder(torch.nn.Module):
"""Class wrapping around a visual feature extractor backbone.
Currently CNN only.
"""
def __init__(
self,
model_name: str = "resnet18",
d_model: int = 512,
in_chans: int = 3,
backend: int = "timm",
**kwargs: Any | None,
):
"""Initialize Visual Encoder.
Args:
model_name (str): Name of the CNN architecture to use (e.g. "resnet18", "resnet50").
d_model (int): Output embedding dimension.
in_chans: the number of input channels of the image.
backend: Which model backend to use. One of {"timm", "torchvision"}
kwargs: see `timm.create_model` and `torchvision.models.resnetX` for kwargs.
"""
super().__init__()
self.model_name = model_name.lower()
self.d_model = d_model
self.backend = backend
if in_chans == 1:
self.in_chans = 3
else:
self.in_chans = in_chans
self.feature_extractor = self.select_feature_extractor(
model_name=self.model_name,
in_chans=self.in_chans,
backend=self.backend,
**kwargs,
)
self.out_layer = torch.nn.Linear(
self.encoder_dim(self.feature_extractor), self.d_model
)
def select_feature_extractor(
self, model_name: str, in_chans: int, backend: str, **kwargs: Any
) -> torch.nn.Module:
"""Select the appropriate feature extractor based on config.
Args:
model_name (str): Name of the CNN architecture to use (e.g. "resnet18", "resnet50").
in_chans: the number of input channels of the image.
backend: Which model backend to use. One of {"timm", "torchvision"}
kwargs: see `timm.create_model` and `torchvision.models.resnetX` for kwargs.
Returns:
a CNN encoder based on the config and backend selected.
"""
if "timm" in backend.lower():
feature_extractor = timm.create_model(
model_name=self.model_name,
in_chans=self.in_chans,
num_classes=0,
**kwargs,
)
elif "torch" in backend.lower():
if model_name.lower() == "resnet18":
feature_extractor = torchvision.models.resnet18(**kwargs)
elif model_name.lower() == "resnet50":
feature_extractor = torchvision.models.resnet50(**kwargs)
else:
raise ValueError(
f"Only `[resnet18, resnet50]` are available when backend is {backend}. Found {model_name}"
)
feature_extractor = torch.nn.Sequential(
*list(feature_extractor.children())[:-1]
)
input_layer = feature_extractor[0]
if in_chans != 3:
feature_extractor[0] = torch.nn.Conv2d(
in_channels=in_chans,
out_channels=input_layer.out_channels,
kernel_size=input_layer.kernel_size,
stride=input_layer.stride,
padding=input_layer.padding,
dilation=input_layer.dilation,
groups=input_layer.groups,
bias=input_layer.bias,
padding_mode=input_layer.padding_mode,
)
else:
raise ValueError(
f"Only ['timm', 'torch'] backends are available! Found {backend}."
)
return feature_extractor
def encoder_dim(self, model: torch.nn.Module) -> int:
"""Compute dummy forward pass of encoder model and get embedding dimension.
Args:
model: a vision encoder model.
Returns:
The embedding dimension size.
"""
_ = model.eval()
dummy_output = model(torch.randn(1, self.in_chans, 224, 224)).squeeze()
_ = model.train() # to be safe
return dummy_output.shape[-1]
def forward(self, img: torch.Tensor) -> torch.Tensor:
"""Forward pass of feature extractor to get feature vector.
Args:
img: Input image tensor of shape (B, C, H, W).
Returns:
feats: Normalized output tensor of shape (B, d_model).
"""
# If grayscale, tile the image to 3 channels.
if img.shape[1] == 1:
img = img.repeat([1, 3, 1, 1]) # (B, nc=3, H, W)
b, c, h, w = img.shape
if c != self.in_chans:
raise ValueError(
f"""Found {c} channels in image but model was configured for {self.in_chans} channels! \n
Hint: have you set the number of anchors in your dataset > 1? \n
If so, make sure to set `in_chans=3 * n_anchors`"""
)
feats = self.feature_extractor(
img
) # (B, out_dim, 1, 1) if using resnet18 backbone.
# Reshape feature vectors
feats = feats.reshape([img.shape[0], -1]) # (B, out_dim)
# Map feature vectors to output dimension using linear layer.
feats = self.out_layer(feats) # (B, d_model)
# Normalize output feature vectors.
feats = F.normalize(feats) # (B, d_model)
return feats
__init__(model_name='resnet18', d_model=512, in_chans=3, backend='timm', **kwargs)
¶
Initialize Visual Encoder.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
model_name
|
str
|
Name of the CNN architecture to use (e.g. "resnet18", "resnet50"). |
'resnet18'
|
d_model
|
int
|
Output embedding dimension. |
512
|
in_chans
|
int
|
the number of input channels of the image. |
3
|
backend
|
int
|
Which model backend to use. One of {"timm", "torchvision"} |
'timm'
|
kwargs
|
Any | None
|
see |
{}
|
Source code in dreem/models/visual_encoder.py
def __init__(
self,
model_name: str = "resnet18",
d_model: int = 512,
in_chans: int = 3,
backend: int = "timm",
**kwargs: Any | None,
):
"""Initialize Visual Encoder.
Args:
model_name (str): Name of the CNN architecture to use (e.g. "resnet18", "resnet50").
d_model (int): Output embedding dimension.
in_chans: the number of input channels of the image.
backend: Which model backend to use. One of {"timm", "torchvision"}
kwargs: see `timm.create_model` and `torchvision.models.resnetX` for kwargs.
"""
super().__init__()
self.model_name = model_name.lower()
self.d_model = d_model
self.backend = backend
if in_chans == 1:
self.in_chans = 3
else:
self.in_chans = in_chans
self.feature_extractor = self.select_feature_extractor(
model_name=self.model_name,
in_chans=self.in_chans,
backend=self.backend,
**kwargs,
)
self.out_layer = torch.nn.Linear(
self.encoder_dim(self.feature_extractor), self.d_model
)
encoder_dim(model)
¶
Compute dummy forward pass of encoder model and get embedding dimension.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
model
|
Module
|
a vision encoder model. |
required |
Returns:
| Type | Description |
|---|---|
int
|
The embedding dimension size. |
Source code in dreem/models/visual_encoder.py
def encoder_dim(self, model: torch.nn.Module) -> int:
"""Compute dummy forward pass of encoder model and get embedding dimension.
Args:
model: a vision encoder model.
Returns:
The embedding dimension size.
"""
_ = model.eval()
dummy_output = model(torch.randn(1, self.in_chans, 224, 224)).squeeze()
_ = model.train() # to be safe
return dummy_output.shape[-1]
forward(img)
¶
Forward pass of feature extractor to get feature vector.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
img
|
Tensor
|
Input image tensor of shape (B, C, H, W). |
required |
Returns:
| Name | Type | Description |
|---|---|---|
feats |
Tensor
|
Normalized output tensor of shape (B, d_model). |
Source code in dreem/models/visual_encoder.py
def forward(self, img: torch.Tensor) -> torch.Tensor:
"""Forward pass of feature extractor to get feature vector.
Args:
img: Input image tensor of shape (B, C, H, W).
Returns:
feats: Normalized output tensor of shape (B, d_model).
"""
# If grayscale, tile the image to 3 channels.
if img.shape[1] == 1:
img = img.repeat([1, 3, 1, 1]) # (B, nc=3, H, W)
b, c, h, w = img.shape
if c != self.in_chans:
raise ValueError(
f"""Found {c} channels in image but model was configured for {self.in_chans} channels! \n
Hint: have you set the number of anchors in your dataset > 1? \n
If so, make sure to set `in_chans=3 * n_anchors`"""
)
feats = self.feature_extractor(
img
) # (B, out_dim, 1, 1) if using resnet18 backbone.
# Reshape feature vectors
feats = feats.reshape([img.shape[0], -1]) # (B, out_dim)
# Map feature vectors to output dimension using linear layer.
feats = self.out_layer(feats) # (B, d_model)
# Normalize output feature vectors.
feats = F.normalize(feats) # (B, d_model)
return feats
select_feature_extractor(model_name, in_chans, backend, **kwargs)
¶
Select the appropriate feature extractor based on config.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
model_name
|
str
|
Name of the CNN architecture to use (e.g. "resnet18", "resnet50"). |
required |
in_chans
|
int
|
the number of input channels of the image. |
required |
backend
|
str
|
Which model backend to use. One of {"timm", "torchvision"} |
required |
kwargs
|
Any
|
see |
{}
|
Returns:
| Type | Description |
|---|---|
Module
|
a CNN encoder based on the config and backend selected. |
Source code in dreem/models/visual_encoder.py
def select_feature_extractor(
self, model_name: str, in_chans: int, backend: str, **kwargs: Any
) -> torch.nn.Module:
"""Select the appropriate feature extractor based on config.
Args:
model_name (str): Name of the CNN architecture to use (e.g. "resnet18", "resnet50").
in_chans: the number of input channels of the image.
backend: Which model backend to use. One of {"timm", "torchvision"}
kwargs: see `timm.create_model` and `torchvision.models.resnetX` for kwargs.
Returns:
a CNN encoder based on the config and backend selected.
"""
if "timm" in backend.lower():
feature_extractor = timm.create_model(
model_name=self.model_name,
in_chans=self.in_chans,
num_classes=0,
**kwargs,
)
elif "torch" in backend.lower():
if model_name.lower() == "resnet18":
feature_extractor = torchvision.models.resnet18(**kwargs)
elif model_name.lower() == "resnet50":
feature_extractor = torchvision.models.resnet50(**kwargs)
else:
raise ValueError(
f"Only `[resnet18, resnet50]` are available when backend is {backend}. Found {model_name}"
)
feature_extractor = torch.nn.Sequential(
*list(feature_extractor.children())[:-1]
)
input_layer = feature_extractor[0]
if in_chans != 3:
feature_extractor[0] = torch.nn.Conv2d(
in_channels=in_chans,
out_channels=input_layer.out_channels,
kernel_size=input_layer.kernel_size,
stride=input_layer.stride,
padding=input_layer.padding,
dilation=input_layer.dilation,
groups=input_layer.groups,
bias=input_layer.bias,
padding_mode=input_layer.padding_mode,
)
else:
raise ValueError(
f"Only ['timm', 'torch'] backends are available! Found {backend}."
)
return feature_extractor
create_visual_encoder(d_model, **encoder_cfg)
¶
Create a visual encoder based on the specified type.
Source code in dreem/models/visual_encoder.py
def create_visual_encoder(d_model: int, **encoder_cfg) -> torch.nn.Module:
"""Create a visual encoder based on the specified type."""
register_encoder("resnet", VisualEncoder)
register_encoder("descriptor", DescriptorVisualEncoder)
# register any custom encoders here
# compatibility with configs that don't specify encoder_type; default to resnet
if not encoder_cfg or "encoder_type" not in encoder_cfg:
encoder_type = "resnet"
return ENCODER_REGISTRY[encoder_type](d_model=d_model, **encoder_cfg)
else:
encoder_type = encoder_cfg.pop("encoder_type")
if encoder_type in ENCODER_REGISTRY:
# choose the relevant encoder configs based on the encoder_type
configs = encoder_cfg[encoder_type]
return ENCODER_REGISTRY[encoder_type](d_model=d_model, **configs)
else:
raise ValueError(
f"Unknown encoder type: {encoder_type}. Please use one of {list(ENCODER_REGISTRY.keys())}"
)
register_encoder(encoder_type, encoder_class)
¶
Register a new encoder type.