Regarding bev_pool_v2

[Unicorncosmos]

Hi ,Iam currently experimenting with the bev_poolv2 and I have implemented the code in python can you guys tell wether it gives the same functionality of the bevpoolv2 please

import torch
from torch import nn
from torch.autograd import Function

class BevPoolV2Function(Function):
    @staticmethod
    def forward(ctx, depth, feat, ranks_depth, ranks_feat, ranks_bev, interval_lengths, interval_starts):
        c = feat.size(4)
        n_intervals = interval_lengths.size(0)

        # Initialize the output tensor
        out = torch.zeros(c, n_intervals, dtype=depth.dtype, device=depth.device)

        # Loop over intervals and calculate the BEV pooling
        for i in range(n_intervals):
            interval_start = interval_starts[i]
            interval_length = interval_lengths[i]
            for cur_c in range(c):
                psum = 0
                for j in range(interval_length):
                    cur_depth = depth[ranks_depth[interval_start + j]]
                    cur_feat = feat[ranks_feat[interval_start + j], :, :, cur_c]
                    psum += cur_feat * cur_depth

                cur_rank = ranks_bev[interval_start]
                out[cur_c, cur_rank] = psum

        ctx.save_for_backward(depth, feat, ranks_depth, ranks_feat, ranks_bev, interval_starts, interval_lengths)

        return out

    @staticmethod
    def backward(ctx, grad_output):
        depth, feat, ranks_depth, ranks_feat, ranks_bev, interval_starts, interval_lengths = ctx.saved_tensors
        c = feat.size(4)
        n_intervals = interval_lengths.size(0)

        # Initialize the gradient tensors
        depth_grad = torch.zeros_like(depth)
        feat_grad = torch.zeros_like(feat)

        # Loop over intervals and calculate gradients
        for i in range(n_intervals):
            interval_start = interval_starts[i]
            interval_length = interval_lengths[i]

            for j in range(interval_length):
                cur_rank = ranks_bev[interval_start + j]
                cur_out_grad = grad_output[:, cur_rank]
                cur_depth = depth[ranks_depth[interval_start + j]]
                cur_feat = feat[ranks_feat[interval_start + j], :, :, :]

                for cur_c in range(c):
                    cur_out_grad_c = cur_out_grad[cur_c]
                    depth_grad[ranks_depth[interval_start + j]] += cur_out_grad_c * cur_feat[cur_c]
                    feat_grad[ranks_feat[interval_start + j], :, :, cur_c] += cur_out_grad_c * cur_depth

        return depth_grad, feat_grad, None, None, None, None, None

class BevPoolV2(nn.Module):
    def __init__(self):
        super(BevPoolV2, self).__init__()

    def forward(self, depth, feat, ranks_depth, ranks_feat, ranks_bev, interval_lengths, interval_starts):
        return BevPoolV2Function.apply(depth, feat, ranks_depth, ranks_feat, ranks_bev, interval_lengths, interval_starts)

[619862306]

Did you verify this code against the results of the cuda implementation?

[Unicorncosmos]

Did you verify this code against the results of the cuda implementation?

yes it is working kindly see my repo where i have implemented it for the Lssviewtransformer and if any changes to be made please open the pr https://github.com/Unicorncosmos/Sfusion_cpubuild/tree/main/CPU_SFUSION_BUILD/LSSViewtransformer