import torch
import torch.nn as nn
from typing import Callable, Optional
[docs]class FlattenTime(nn.Flatten):
"""
Utility layer which always flattens first two dimensions. Meant
to convert a tensor of dimensions (Batch, Time, Channels, Height, Width)
into a tensor of (Batch*Time, Channels, Height, Width).
"""
def __init__(self):
super().__init__(start_dim=0, end_dim=1)
[docs]class UnflattenTime(nn.Module):
"""
Utility layer which always unflattens (expands) the first dimension into two separate ones.
Meant to convert a tensor of dimensions (Batch*Time, Channels, Height, Width)
into a tensor of (Batch, Time, Channels, Height, Width).
"""
def __init__(self, batch_size: int):
super().__init__()
self.batch_size = batch_size
[docs] def forward(self, x):
num_time_steps = x.shape[0] // self.batch_size
return x.unflatten(0, (self.batch_size, num_time_steps))
[docs]class SqueezeMixin:
"""
Utility mixin class that will wrap the __init__ and forward call of other classes.
The wrapped __init__ will provide two additional parameters batch_size and num_timesteps
and the wrapped forward will unpack and repack the first dimension into batch and time.
"""
def squeeze_init(self, batch_size: Optional[int], num_timesteps: Optional[int]):
if not batch_size and not num_timesteps:
raise TypeError("You need to specify either batch_size or num_timesteps.")
if not batch_size:
batch_size = -1
if not num_timesteps:
num_timesteps = -1
self.batch_size = int(batch_size)
self.num_timesteps = int(num_timesteps)
def squeeze_forward(self, input_data: torch.Tensor, forward_method: Callable):
inflated_input = input_data.reshape(
self.batch_size, self.num_timesteps, *input_data.shape[1:]
)
inflated_output = forward_method(inflated_input)
return inflated_output.flatten(start_dim=0, end_dim=1)
def squeeze_param_dict(self, param_dict: dict) -> dict:
param_dict.update(
batch_size=self.batch_size,
num_timesteps=self.num_timesteps,
)
return param_dict
