File size: 4,453 Bytes
6a176cb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 | 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}')
|