pointcept/models/stratified_transformer/stratified_transformer_v1m2_refine.py

"""
Stratified Transformer

Modified from https://github.com/dvlab-research/Stratified-Transformer

Author: Xiaoyang Wu (xiaoyang.wu.cs@gmail.com)
Please cite our work if the code is helpful to you.
"""

from copy import deepcopy
import torch
import torch.nn as nn

try:
    import torch_points_kernels as tp
except ImportError:
    tp = None

try:
    from torch_points3d.modules.KPConv.kernels import KPConvLayer
    from torch_points3d.core.common_modules import FastBatchNorm1d
except ImportError:
    KPConvLayer = None
    FastBatchNorm1d = None

from torch_scatter import scatter_softmax
from timm.layers import DropPath, trunc_normal_
from torch_geometric.nn.pool import voxel_grid

try:
    import pointops2.pointops as pointops
except ImportError:
    pointops = None

from pointcept.models.builder import MODELS


def offset2batch(offset):
    return (
        torch.cat(
            [
                (
                    torch.tensor([i] * (o - offset[i - 1]))
                    if i > 0
                    else torch.tensor([i] * o)
                )
                for i, o in enumerate(offset)
            ],
            dim=0,
        )
        .long()
        .to(offset.device)
    )


def grid_sample(coords, batch, size, start, return_p2v=True):
    cluster = voxel_grid(coords, batch, size, start=start)

    if not return_p2v:
        unique, cluster = torch.unique(cluster, sorted=True, return_inverse=True)
        return cluster
    else:
        unique, cluster, counts = torch.unique(
            cluster, sorted=True, return_inverse=True, return_counts=True
        )

        # obtain p2v_map
        n = unique.shape[0]
        k = counts.max().item()
        p2v_map = cluster.new_zeros(n, k)
        mask = torch.arange(k).cuda().unsqueeze(0) < counts.unsqueeze(-1)
        p2v_map[mask] = torch.argsort(cluster)
        return cluster, p2v_map, counts


class WindowAttention(nn.Module):
    """Window based multi-head self attention (W-MSA) module with relative position bias.
    It supports both of shifted and non-shifted window.
    """

    def __init__(
        self,
        embed_channels,
        num_heads,
        window_size,
        quant_size,
        attn_drop=0.0,
        proj_drop=0.0,
        scale=None,
        rel_query=True,
        rel_key=True,
        rel_value=True,
        qkv_bias=True,
    ):
        super().__init__()
        self.embed_channels = embed_channels
        self.head_channels = embed_channels // num_heads
        self.num_heads = num_heads
        self.scale = scale or self.head_channels**-0.5

        self.window_size = window_size
        self.quant_size = quant_size

        self.rel_query = rel_query
        self.rel_key = rel_key
        self.rel_value = rel_value

        self.quant_grid_length = int((2 * window_size + 1e-4) // quant_size)

        assert self.rel_query and self.rel_key
        if rel_query:
            self.relative_pos_query_table = nn.Parameter(
                torch.zeros(
                    2 * self.quant_grid_length, self.num_heads, self.head_channels, 3
                )
            )
            trunc_normal_(self.relative_pos_query_table, std=0.02)

        if rel_key:
            self.relative_pos_key_table = nn.Parameter(
                torch.zeros(
                    2 * self.quant_grid_length, self.num_heads, self.head_channels, 3
                )
            )
            trunc_normal_(self.relative_pos_query_table, std=0.02)

        if rel_value:
            self.relative_pos_value_table = nn.Parameter(
                torch.zeros(
                    2 * self.quant_grid_length, self.num_heads, self.head_channels, 3
                )
            )
            trunc_normal_(self.relative_pos_query_table, std=0.02)

        self.qkv = nn.Linear(embed_channels, embed_channels * 3, bias=qkv_bias)
        self.attn_drop = nn.Dropout(attn_drop, inplace=True)
        self.proj = nn.Linear(embed_channels, embed_channels)
        self.proj_drop = nn.Dropout(proj_drop, inplace=True)

        self.softmax = nn.Softmax(dim=-1)

    def forward(self, feats, coords, index_0, index_1, index_0_offsets, n_max):
        n, c = feats.shape
        m = index_0.shape[0]

        assert index_0.shape[0] == index_1.shape[0]

        qkv = (
            self.qkv(feats)
            .reshape(n, 3, self.num_heads, c // self.num_heads)
            .permute(1, 0, 2, 3)
            .contiguous()
        )
        query, key, value = qkv[0], qkv[1], qkv[2]
        query = query * self.scale
        attn_flat = pointops.attention_step1_v2(
            query.float(), key.float(), index_1.int(), index_0_offsets.int(), n_max
        )

        # Position embedding
        relative_position = coords[index_0] - coords[index_1]
        relative_position = torch.round(relative_position * 100000) / 100000
        relative_position_index = torch.div(
            relative_position + 2 * self.window_size - 1e-4,
            self.quant_size,
            rounding_mode="trunc",
        )
        # relative_position_index = (relative_position + 2 * self.window_size - 1e-4) // self.quant_size
        assert (relative_position_index >= 0).all()
        assert (relative_position_index <= 2 * self.quant_grid_length - 1).all()

        if self.rel_query and self.rel_key:
            relative_position_bias = pointops.dot_prod_with_idx_v3(
                query.float(),
                index_0_offsets.int(),
                n_max,
                key.float(),
                index_1.int(),
                self.relative_pos_query_table.float(),
                self.relative_pos_key_table.float(),
                relative_position_index.int(),
            )
        elif self.rel_query:
            relative_position_bias = pointops.dot_prod_with_idx(
                query.float(),
                index_0.int(),
                self.relative_pos_query_table.float(),
                relative_position_index.int(),
            )  # [M, num_heads]
        elif self.rel_key:
            relative_position_bias = pointops.dot_prod_with_idx(
                key.float(),
                index_1.int(),
                self.relative_pos_key_table.float(),
                relative_position_index.int(),
            )  # [M, num_heads]
        else:
            relative_position_bias = 0.0

        attn_flat += relative_position_bias
        softmax_attn_flat = scatter_softmax(src=attn_flat, index=index_0, dim=0)

        if self.rel_value:
            x = pointops.attention_step2_with_rel_pos_value_v2(
                softmax_attn_flat.float(),
                value.float(),
                index_0_offsets.int(),
                n_max,
                index_1.int(),
                self.relative_pos_value_table.float(),
                relative_position_index.int(),
            )
        else:
            x = pointops.attention_step2(
                softmax_attn_flat.float(), value.float(), index_0.int(), index_1.int()
            )

        x = x.view(n, c)
        x = self.proj(x)
        x = self.proj_drop(x)
        return x


class MLP(nn.Module):
    def __init__(self, in_channels, hidden_channels=None, out_channels=None, drop=0.0):
        super().__init__()
        out_channels = out_channels or in_channels
        hidden_channels = hidden_channels or in_channels
        self.fc1 = nn.Linear(in_channels, hidden_channels)
        self.act = nn.GELU()
        self.fc2 = nn.Linear(hidden_channels, out_channels)
        self.drop = nn.Dropout(drop, inplace=True)

    def forward(self, x):
        x = self.fc1(x)
        x = self.act(x)
        x = self.drop(x)
        x = self.fc2(x)
        x = self.drop(x)
        return x


class Block(nn.Module):
    def __init__(
        self,
        embed_channels,
        num_heads,
        window_size,
        quant_size,
        mlp_expend_ratio=4.0,
        drop_path=0.0,
        qk_scale=None,
        rel_query=True,
        rel_key=True,
        rel_value=True,
        qkv_bias=True,
    ):
        super().__init__()
        self.norm1 = nn.LayerNorm(embed_channels)
        self.attn = WindowAttention(
            embed_channels,
            num_heads,
            window_size,
            quant_size,
            scale=qk_scale,
            rel_query=rel_query,
            rel_key=rel_key,
            rel_value=rel_value,
            qkv_bias=qkv_bias,
        )
        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
        self.norm2 = nn.LayerNorm(embed_channels)
        self.mlp = MLP(
            in_channels=embed_channels,
            hidden_channels=int(embed_channels * mlp_expend_ratio),
        )

    def forward(self, feats, coords, index_0, index_1, index_0_offsets, n_max):
        short_cut = feats
        feats = self.norm1(feats)
        feats = self.attn(feats, coords, index_0, index_1, index_0_offsets, n_max)

        feats = short_cut + self.drop_path(feats)
        feats += self.drop_path(self.mlp(self.norm2(feats)))
        return feats


class BasicLayer(nn.Module):
    def __init__(
        self,
        embed_channels,
        out_channels,
        depth,
        num_heads,
        window_size,
        quant_size,
        mlp_expend_ratio=4.0,
        down_ratio=0.25,
        down_num_sample=16,
        drop_path=None,
        qk_scale=None,
        down=True,
        rel_query=True,
        rel_key=True,
        rel_value=True,
        qkv_bias=True,
    ):
        super().__init__()
        self.depth = depth
        self.window_size = window_size
        self.quant_size = quant_size
        self.down_ratio = down_ratio

        if isinstance(drop_path, list):
            drop_path = drop_path
            assert len(drop_path) == depth
        elif isinstance(drop_path, float):
            drop_path = [deepcopy(drop_path) for _ in range(depth)]
        else:
            drop_path = [0.0 for _ in range(depth)]

        self.blocks = nn.ModuleList()
        for i in range(depth):
            block = Block(
                embed_channels,
                num_heads,
                window_size,
                quant_size,
                mlp_expend_ratio=mlp_expend_ratio,
                drop_path=drop_path[i],
                qk_scale=qk_scale,
                rel_query=rel_query,
                rel_key=rel_key,
                rel_value=rel_value,
                qkv_bias=qkv_bias,
            )
            self.blocks.append(block)

        self.down = (
            TransitionDown(embed_channels, out_channels, down_ratio, down_num_sample)
            if down
            else None
        )

    def forward(self, feats, coords, offset):
        # window_size -> [window_size, window_size, window_size]
        window_size = torch.tensor(
            [self.window_size] * 3, dtype=coords.dtype, device=coords.device
        )
        new_window_size = 2 * torch.tensor(
            [self.window_size] * 3, dtype=coords.dtype, device=coords.device
        )
        batch = offset2batch(offset)

        # compute new offset
        new_offset = [int(offset[0].item() * self.down_ratio) + 1]
        count = int(offset[0].item() * self.down_ratio) + 1
        for i in range(1, offset.shape[0]):
            count += (
                int((offset[i].item() - offset[i - 1].item()) * self.down_ratio) + 1
            )
            new_offset.append(count)
        new_offset = torch.cuda.IntTensor(new_offset)
        down_idx = pointops.furthestsampling(coords, offset.int(), new_offset.int())

        # compute window mapping
        coords_min = coords.min(0).values
        v2p_map, p2v_map, counts = grid_sample(coords, batch, window_size, start=None)
        shift_size = window_size * 1 / 2
        shift_v2p_map, shift_p2v_map, shift_counts = grid_sample(
            coords + shift_size, batch, window_size, start=coords_min
        )

        new_v2p_map, new_p2v_map, new_counts = grid_sample(
            coords, batch, new_window_size, start=None
        )
        shift_size = new_window_size * 1 / 2
        shift_new_v2p_map, shift_new_p2v_map, shift_new_counts = grid_sample(
            coords + shift_size, batch, new_window_size, start=coords_min
        )

        # stratified attention
        for i, blk in enumerate(self.blocks):
            p2v_map_blk = p2v_map if i % 2 == 0 else shift_p2v_map
            counts_blk = counts if i % 2 == 0 else shift_counts

            new_p2v_map_blk = new_p2v_map if i % 2 == 0 else shift_new_p2v_map
            new_counts_blk = new_counts if i % 2 == 0 else shift_new_counts

            n, k = p2v_map_blk.shape
            mask = torch.arange(k).unsqueeze(0).cuda() < counts_blk.unsqueeze(-1)
            mask_mat = mask.unsqueeze(-1) & mask.unsqueeze(-2)
            index_0 = p2v_map_blk.unsqueeze(-1).expand(-1, -1, k)[mask_mat]
            index_1 = p2v_map_blk.unsqueeze(1).expand(-1, k, -1)[mask_mat]

            down_mask = torch.zeros_like(batch).bool()
            down_mask[down_idx.long()] = True
            down_mask = down_mask[new_p2v_map_blk]  # [n, k], down sample mask
            n, k = new_p2v_map_blk.shape
            mask = torch.arange(k).unsqueeze(0).cuda() < new_counts_blk.unsqueeze(
                -1
            )  # [n, k]
            down_mask = down_mask & mask  # down sample and window mask
            # [n, k, k] query: dense point in large windows; key: sparse point in large windows
            mask_mat = mask.unsqueeze(-1) & down_mask.unsqueeze(-2)

            if i % 2 == 0:
                # [n, k, 3]
                # window_coord = (coords[new_p2v_map_blk] - coords_min) // window_size
                window_coord = torch.div(
                    coords[new_p2v_map_blk] - coords_min,
                    window_size,
                    rounding_mode="trunc",
                )
            else:
                # [n, k, 3]
                # window_coord = (coords[new_p2v_map_blk] - coords_min + 1/2 * window_size) // window_size
                window_coord = torch.div(
                    coords[new_p2v_map_blk] - coords_min + 1 / 2 * window_size,
                    window_size,
                    rounding_mode="trunc",
                )
            # [n, k, k], whether pair points are in same small windows
            mask_mat_prev = (
                window_coord.unsqueeze(2) != window_coord.unsqueeze(1)
            ).any(-1)
            mask_mat = mask_mat & mask_mat_prev

            new_index_0 = new_p2v_map_blk.unsqueeze(-1).expand(-1, -1, k)[mask_mat]
            new_index_1 = new_p2v_map_blk.unsqueeze(1).expand(-1, k, -1)[mask_mat]

            index_0 = torch.cat([index_0, new_index_0], 0)
            index_1 = torch.cat([index_1, new_index_1], 0)

            # rearrange index for acceleration
            index_0, indices = torch.sort(index_0)
            index_1 = index_1[indices]
            index_0_counts = index_0.bincount()
            n_max = index_0_counts.max()
            index_0_offsets = index_0_counts.cumsum(dim=-1)
            index_0_offsets = torch.cat(
                [torch.zeros(1, dtype=torch.long).cuda(), index_0_offsets], 0
            )

            feats = blk(feats, coords, index_0, index_1, index_0_offsets, n_max)

        if self.down:
            feats_down, coords_down, offset_down = self.down(feats, coords, offset)
        else:
            feats_down, coords_down, offset_down = None, None, None

        return feats, coords, offset, feats_down, coords_down, offset_down


class TransitionDown(nn.Module):
    def __init__(self, in_channels, out_channels, ratio, k, norm_layer=nn.LayerNorm):
        super().__init__()
        self.ratio = ratio
        self.k = k
        self.norm = norm_layer(in_channels) if norm_layer else None
        self.linear = nn.Linear(in_channels, out_channels, bias=False)
        self.pool = nn.MaxPool1d(k)

    def forward(self, feats, coords, offset):
        new_offset, count = [int(offset[0].item() * self.ratio) + 1], int(
            offset[0].item() * self.ratio
        ) + 1
        for i in range(1, offset.shape[0]):
            count += ((offset[i].item() - offset[i - 1].item()) * self.ratio) + 1
            new_offset.append(count)
        new_offset = torch.cuda.IntTensor(new_offset)
        idx = pointops.furthestsampling(coords, offset, new_offset)  # (m)
        new_coords = coords[idx.long(), :]  # (m, 3)

        feats = pointops.queryandgroup(
            self.k, coords, new_coords, feats, None, offset, new_offset, use_xyz=False
        )  # (m, nsample, 3+c)
        m, k, c = feats.shape
        feats = (
            self.linear(self.norm(feats.view(m * k, c)).view(m, k, c))
            .transpose(1, 2)
            .contiguous()
        )
        feats = self.pool(feats).squeeze(-1)  # (m, c)
        return feats, new_coords, new_offset


class TransitionUp(nn.Module):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels

        self.linear1 = nn.Sequential(
            nn.LayerNorm(out_channels), nn.Linear(out_channels, out_channels)
        )

        self.linear2 = nn.Sequential(
            nn.LayerNorm(in_channels), nn.Linear(in_channels, out_channels)
        )

    def forward(self, feats, coords, offset, skip_feats, skip_coords, skip_offset):
        feats = self.linear1(skip_feats) + pointops.interpolation(
            coords, skip_coords, self.linear2(feats), offset, skip_offset
        )
        return feats, skip_coords, skip_offset


class KPConvSimpleBlock(nn.Module):
    def __init__(
        self,
        in_channels,
        out_channels,
        prev_grid_size,
        sigma=1.0,
        negative_slope=0.2,
        bn_momentum=0.02,
    ):
        super().__init__()
        self.kpconv = KPConvLayer(
            in_channels,
            out_channels,
            point_influence=prev_grid_size * sigma,
            add_one=False,
        )
        self.bn = FastBatchNorm1d(out_channels, momentum=bn_momentum)
        self.activation = nn.LeakyReLU(negative_slope=negative_slope)

    def forward(self, feats, xyz, batch, neighbor_idx):
        # feats: [N, C]
        # coords: [N, 3]
        # batch: [N,]
        # neighbor_idx: [N, M]

        feats = self.kpconv(xyz, xyz, neighbor_idx, feats)
        feats = self.activation(self.bn(feats))
        return feats


class KPConvResBlock(nn.Module):
    def __init__(
        self,
        in_channels,
        out_channels,
        prev_grid_size,
        sigma=1.0,
        negative_slope=0.2,
        bn_momentum=0.02,
    ):
        super().__init__()
        d_2 = out_channels // 4
        activation = nn.LeakyReLU(negative_slope=negative_slope)
        self.unary_1 = torch.nn.Sequential(
            nn.Linear(in_channels, d_2, bias=False),
            FastBatchNorm1d(d_2, momentum=bn_momentum),
            activation,
        )
        self.unary_2 = torch.nn.Sequential(
            nn.Linear(d_2, out_channels, bias=False),
            FastBatchNorm1d(out_channels, momentum=bn_momentum),
            activation,
        )
        self.kpconv = KPConvLayer(
            d_2, d_2, point_influence=prev_grid_size * sigma, add_one=False
        )
        self.bn = FastBatchNorm1d(out_channels, momentum=bn_momentum)
        self.activation = activation

        if in_channels != out_channels:
            self.shortcut_op = torch.nn.Sequential(
                nn.Linear(in_channels, out_channels, bias=False),
                FastBatchNorm1d(out_channels, momentum=bn_momentum),
            )
        else:
            self.shortcut_op = nn.Identity()

    def forward(self, feats, xyz, batch, neighbor_idx):
        # feats: [N, C]
        # coords: [N, 3]
        # batch: [N,]
        # neighbor_idx: [N, M]

        shortcut = feats
        feats = self.unary_1(feats)
        feats = self.kpconv(xyz, xyz, neighbor_idx, feats)
        feats = self.unary_2(feats)
        shortcut = self.shortcut_op(shortcut)
        feats += shortcut
        return feats


@MODELS.register_module("ST-v1m2")
class StratifiedTransformer(nn.Module):
    def __init__(
        self,
        in_channels,
        num_classes,
        channels=(48, 96, 192, 384, 384),
        num_heads=(6, 12, 24, 24),
        depths=(3, 9, 3, 3),
        window_size=(0.2, 0.4, 0.8, 1.6),
        quant_size=(0.01, 0.02, 0.04, 0.08),
        mlp_expend_ratio=4.0,
        down_ratio=0.25,
        down_num_sample=16,
        kp_ball_radius=2.5 * 0.02,
        kp_max_neighbor=34,
        kp_grid_size=0.02,
        kp_sigma=1.0,
        drop_path_rate=0.2,
        rel_query=True,
        rel_key=True,
        rel_value=True,
        qkv_bias=True,
        stem=True,
    ):
        super().__init__()
        assert (
            KPConvLayer is not None and FastBatchNorm1d is not None
        ), "Please make sure torch_points3d is installed"
        assert tp is not None, "Please make sure torch_points_kernels is installed"
        assert pointops is not None, "Please make sure pointops2 is installed"
        # stochastic depth decay rule
        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))]
        self.kp_ball_radius = kp_ball_radius
        self.kp_max_neighbor = kp_max_neighbor
        self.stem = stem
        if stem:
            self.point_embed = nn.ModuleList(
                [
                    KPConvSimpleBlock(
                        in_channels, channels[0], kp_grid_size, sigma=kp_sigma
                    ),
                    KPConvResBlock(
                        channels[0], channels[0], kp_grid_size, sigma=kp_sigma
                    ),
                ]
            )
            self.down = TransitionDown(
                channels[0], channels[1], down_ratio, down_num_sample
            )
        else:
            assert channels[0] == channels[1]
            self.point_embed = nn.ModuleList(
                [
                    KPConvSimpleBlock(
                        in_channels, channels[1], kp_grid_size, sigma=kp_sigma
                    ),
                ]
            )

        num_layers = len(depths)
        self.layers = nn.ModuleList()
        for i in range(num_layers):
            layer = BasicLayer(
                embed_channels=channels[i + 1],
                out_channels=channels[i + 2] if i < num_layers - 1 else channels[i + 1],
                depth=depths[i],
                num_heads=num_heads[i],
                window_size=window_size[i],
                quant_size=quant_size[i],
                mlp_expend_ratio=mlp_expend_ratio,
                down_ratio=down_ratio,
                down_num_sample=down_num_sample,
                drop_path=dpr[sum(depths[:i]) : sum(depths[: i + 1])],
                rel_query=rel_query,
                rel_key=rel_key,
                rel_value=rel_value,
                qkv_bias=qkv_bias,
                down=True if i < num_layers - 1 else False,
            )
            self.layers.append(layer)

        self.up = nn.ModuleList(
            [
                TransitionUp(channels[i + 1], channels[i])
                for i in reversed(range(1, num_layers))
            ]
        )
        if self.stem:
            self.up.append(TransitionUp(channels[1], channels[0]))

        self.classifier = nn.Sequential(
            nn.Linear(channels[0], channels[0]),
            nn.BatchNorm1d(channels[0]),
            nn.ReLU(inplace=True),
            nn.Linear(channels[0], num_classes),
        )

        self.init_weights()

    def forward(self, data_dict):
        feats = data_dict["feat"]
        coords = data_dict["coord"]
        offset = data_dict["offset"].int()
        batch = offset2batch(offset)
        neighbor_idx = tp.ball_query(
            self.kp_ball_radius,
            self.kp_max_neighbor,
            coords,
            coords,
            mode="partial_dense",
            batch_x=batch,
            batch_y=batch,
        )[0]

        feats_stack = []
        coords_stack = []
        offset_stack = []

        for i, layer in enumerate(self.point_embed):
            feats = layer(feats, coords, batch, neighbor_idx)

        feats = feats.contiguous()
        if self.stem:
            feats_stack.append(feats)
            coords_stack.append(coords)
            offset_stack.append(offset)
            feats, coords, offset = self.down(feats, coords, offset)

        for i, layer in enumerate(self.layers):
            feats, coords, offset, feats_down, coords_down, offset_down = layer(
                feats, coords, offset
            )

            feats_stack.append(feats)
            coords_stack.append(coords)
            offset_stack.append(offset)

            feats = feats_down
            coords = coords_down
            offset = offset_down

        feats = feats_stack.pop()
        coords = coords_stack.pop()
        offset = offset_stack.pop()

        for i, up in enumerate(self.up):
            feats, coords, offset = up(
                feats,
                coords,
                offset,
                feats_stack.pop(),
                coords_stack.pop(),
                offset_stack.pop(),
            )

        out = self.classifier(feats)
        return out

    def init_weights(self):
        """Initialize the weights in backbone."""

        def _init_weights(m):
            if isinstance(m, nn.Linear):
                trunc_normal_(m.weight, std=0.02)
                if isinstance(m, nn.Linear) and m.bias is not None:
                    nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.LayerNorm) or isinstance(m, nn.BatchNorm1d):
                nn.init.constant_(m.bias, 0)
                nn.init.constant_(m.weight, 1.0)

        self.apply(_init_weights)