| import torch |
| from torch import nn |
|
|
|
|
| class ReuseUpdater: |
| def __call__(self, previous_memory, *args, **kwargs): |
| return previous_memory |
|
|
|
|
| class GatedResidualMemoryUpdater(nn.Module): |
| def __init__( |
| self, |
| channels, |
| ego_dim=0, |
| feature_format='spatial', |
| residual_scale_init=0.0, |
| ): |
| super().__init__() |
| if feature_format not in ('spatial', 'tokens'): |
| raise ValueError("feature_format must be 'spatial' or 'tokens'") |
| self.channels = channels |
| self.ego_dim = ego_dim |
| self.feature_format = feature_format |
| in_channels = channels * 2 + ego_dim |
| if feature_format == 'spatial': |
| self.gate = nn.Conv2d(in_channels, channels, kernel_size=1) |
| self.residual = nn.Conv2d(in_channels, channels, kernel_size=1) |
| else: |
| self.gate = nn.Linear(in_channels, channels) |
| self.residual = nn.Linear(in_channels, channels) |
| self.residual_scale = nn.Parameter(torch.tensor(float(residual_scale_init))) |
|
|
| def forward(self, previous_memory, partial_current=None, ego_motion_embedding=None): |
| if isinstance(previous_memory, list): |
| return [ |
| self.forward(prev, part, ego_motion_embedding) |
| for prev, part in zip(previous_memory, partial_current) |
| ] |
| if isinstance(previous_memory, tuple): |
| return tuple( |
| self.forward(prev, part, ego_motion_embedding) |
| for prev, part in zip(previous_memory, partial_current) |
| ) |
| if partial_current is None: |
| raise ValueError('partial_current is required for gated residual updates') |
| if previous_memory.shape != partial_current.shape: |
| raise ValueError( |
| f'partial_current shape {tuple(partial_current.shape)} must match ' |
| f'previous_memory shape {tuple(previous_memory.shape)}') |
|
|
| x = self._concat_inputs(previous_memory, partial_current, ego_motion_embedding) |
| gate = torch.sigmoid(self.gate(x)) |
| residual = self.residual(x) |
| if self.feature_format == 'spatial': |
| gate = self._from_spatial(gate, previous_memory) |
| residual = self._from_spatial(residual, previous_memory) |
| return previous_memory + self.residual_scale * gate * residual |
|
|
| def _concat_inputs(self, previous_memory, partial_current, ego_motion_embedding): |
| if self.feature_format == 'spatial': |
| if previous_memory.dim() != 5: |
| raise ValueError('spatial memory must have shape B,N_cam,C,H,W') |
| pieces = [previous_memory, partial_current] |
| if ego_motion_embedding is not None: |
| pieces.append(self._expand_spatial_ego(ego_motion_embedding, previous_memory)) |
| return self._to_spatial(torch.cat(pieces, dim=2)) |
|
|
| if previous_memory.dim() != 4: |
| raise ValueError('token memory must have shape B,N_cam,N_token,C') |
| pieces = [previous_memory, partial_current] |
| if ego_motion_embedding is not None: |
| pieces.append(self._expand_token_ego(ego_motion_embedding, previous_memory)) |
| return torch.cat(pieces, dim=-1) |
|
|
| @staticmethod |
| def _to_spatial(x): |
| b, n_cam, channels, h, w = x.shape |
| return x.reshape(b * n_cam, channels, h, w) |
|
|
| @staticmethod |
| def _from_spatial(x, reference): |
| b, n_cam, _, h, w = reference.shape |
| return x.reshape(b, n_cam, -1, h, w) |
|
|
| @staticmethod |
| def _expand_spatial_ego(ego_motion_embedding, reference): |
| b, n_cam, _, h, w = reference.shape |
| return ego_motion_embedding[:, None, :, None, None].expand(b, n_cam, -1, h, w) |
|
|
| @staticmethod |
| def _expand_token_ego(ego_motion_embedding, reference): |
| b, n_cam, n_token, _ = reference.shape |
| return ego_motion_embedding[:, None, None, :].expand(b, n_cam, n_token, -1) |
|
|
|
|
| def build_memory_updater(spec, **kwargs): |
| if callable(spec) and not isinstance(spec, str): |
| return spec |
| if isinstance(spec, dict): |
| spec = spec.copy() |
| name = spec.pop('type', spec.pop('name', 'reuse')) |
| spec.update(kwargs) |
| else: |
| name = spec |
| spec = kwargs |
|
|
| if name == 'reuse': |
| return ReuseUpdater() |
| if name == 'gated_residual': |
| return GatedResidualMemoryUpdater(**spec) |
| raise KeyError(f'Unsupported temporal memory updater: {name}') |
|
|
