mmagic.models.archs.ensemble

Module Contents

Classes

SpatialTemporalEnsemble

Apply spatial and temporal ensemble and compute outputs.

class mmagic.models.archs.ensemble.SpatialTemporalEnsemble(is_temporal_ensemble: Optional[bool] = False)

Bases: torch.nn.Module

Apply spatial and temporal ensemble and compute outputs.

Parameters

is_temporal_ensemble (bool, optional) – Whether to apply ensemble temporally. If True, the sequence will also be flipped temporally. If the input is an image, this argument must be set to False. Default: False.

_transform(imgs: torch.Tensor, mode: str) → torch.Tensor

Apply spatial transform (flip, rotate) to the images.

Parameters

• imgs (torch.Tensor) – The images to be transformed/

• mode (str) – The mode of transform. Supported values are ‘vertical’, ‘horizontal’, and ‘transpose’, corresponding to vertical flip, horizontal flip, and rotation, respectively.

Returns

Output of the model with spatial ensemble applied.

Return type

torch.Tensor

spatial_ensemble(imgs: torch.Tensor, model: torch.nn.Module) → torch.Tensor

Apply spatial ensemble.

Parameters

• imgs (torch.Tensor) – The images to be processed by the model. Its size should be either (n, t, c, h, w) or (n, c, h, w).

• model (nn.Module) – The model to process the images.

Returns

Output of the model with spatial ensemble applied.

Return type

torch.Tensor

forward(imgs: torch.Tensor, model: torch.nn.Module) → torch.Tensor

Apply spatial and temporal ensemble.

Parameters

• imgs (torch.Tensor) – The images to be processed by the model. Its size should be either (n, t, c, h, w) or (n, c, h, w).

• model (nn.Module) – The model to process the images.

Returns

Output of the model with spatial ensemble applied.

Return type

torch.Tensor