from __future__ import annotations
from typing import Literal
import torch
from torch import nn
from torch_lattice import SparseTensor
from torch_lattice.nn import functional as F
__all__ = [
"AvgPool3d",
"GlobalAvgPool",
"GlobalMaxPool",
"MaxPool3d",
"Pool3d",
"SumPool3d",
]
[docs]
class Pool3d(nn.Module):
"""Local sparse 3D pooling over generated output support."""
def __init__(
self,
*,
mode: Literal["sum", "max", "avg"],
kernel_size=2,
stride=2,
padding=0,
dilation=1,
) -> None:
super().__init__()
self.mode = mode
self.kernel_size = kernel_size
self.stride = stride
self.padding = padding
self.dilation = dilation
[docs]
def forward(self, input: SparseTensor) -> SparseTensor:
return F.pool3d(
input,
mode=self.mode,
kernel_size=self.kernel_size,
stride=self.stride,
padding=self.padding,
dilation=self.dilation,
)
[docs]
class SumPool3d(Pool3d):
def __init__(self, kernel_size=2, stride=2, padding=0, dilation=1) -> None:
super().__init__(
mode="sum",
kernel_size=kernel_size,
stride=stride,
padding=padding,
dilation=dilation,
)
[docs]
class MaxPool3d(Pool3d):
def __init__(self, kernel_size=2, stride=2, padding=0, dilation=1) -> None:
super().__init__(
mode="max",
kernel_size=kernel_size,
stride=stride,
padding=padding,
dilation=dilation,
)
[docs]
class AvgPool3d(Pool3d):
def __init__(self, kernel_size=2, stride=2, padding=0, dilation=1) -> None:
super().__init__(
mode="avg",
kernel_size=kernel_size,
stride=stride,
padding=padding,
dilation=dilation,
)
[docs]
class GlobalAvgPool(nn.Module):
[docs]
def forward(self, input: SparseTensor) -> torch.Tensor:
return F.global_avg_pool(input)
[docs]
class GlobalMaxPool(nn.Module):
[docs]
def forward(self, input: SparseTensor) -> torch.Tensor:
return F.global_max_pool(input)