<filename>aidet/model/head/gfl_head.py<gh_stars>0
import math
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.distributed as dist


def reduce_mean(tensor):
    if not (dist.is_available() and dist.is_initialized()):
        return tensor
    tensor = tensor.clone()
    dist.all_reduce(tensor.true_divide(dist.get_world_size()), op=dist.ReduceOp.SUM)
    return tensor

class Integral(nn.Module):
    """
    A fixed layer for calculating integral result from distribution.
    This layer calculates the target location by:math: 'sum{P(y_i) * y_i}',
    P(y_i) denotes the softmax vector that represents the discrete distribution.
    y_i denotes the discrete set, usually {0, 1, 2,..., reg_max}
    Args:
        reg_max (int): the maximal value of the discrete set. Default:16.
            you may reset it according to your new dataset or related settings.
    """
    def __init__(self, reg_max=16):
        super(Integral, self).__init__()
        self.reg_max = reg_max
        self.register_buffer(
            "project", torch.linspace(0, self.reg_max, self.reg_max + 1)
        )

    def forward(self, x):
        """
        Forward feature from the regression head to get integral result of bounding box location.
        Args:
            :param x (Tensor): Features of the regression head, shape (N, 4*(n+1)), n is self.reg_max.
            :return x (Tensor): Integral result of box location, i.e., distance offsets from box center in four directions, shape (N, 4).
        """
        shape = x.size()  # x.shape [b, 2125, 32]
        x = F.softmax(x.reshape(*shape[:-1], 4, self.reg_max + 1), dim=-1)  # [b, 2125, 4, 8]
        x = F.linear(x, self.project.type_as(x)).reshape(*shape[:-1], 4)
        return x
<filename>aidet/model/head/gfl_head.py<gh_stars>0
import math
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.distributed as dist


def reduce_mean(tensor):
    if not (dist.is_available() and dist.is_initialized()):
        return tensor
    tensor = tensor.clone()
    dist.all_reduce(tensor.true_divide(dist.get_world_size()), op=dist.ReduceOp.SUM)
    return tensor

class Integral(nn.Module):
    """
    A fixed layer for calculating integral result from distribution.
    This layer calculates the target location by:math: 'sum{P(y_i) * y_i}',
    P(y_i) denotes the softmax vector that represents the discrete distribution.
    y_i denotes the discrete set, usually {0, 1, 2,..., reg_max}
    Args:
        reg_max (int): the maximal value of the discrete set. Default:16.
            you may reset it according to your new dataset or related settings.
    """
    def __init__(self, reg_max=16):
        super(Integral, self).__init__()
        self.reg_max = reg_max
        self.register_buffer(
            "project", torch.linspace(0, self.reg_max, self.reg_max + 1)
        )

    def forward(self, x):
        """
        Forward feature from the regression head to get integral result of bounding box location.
        Args:
            :param x (Tensor): Features of the regression head, shape (N, 4*(n+1)), n is self.reg_max.
            :return x (Tensor): Integral result of box location, i.e., distance offsets from box center in four directions, shape (N, 4).
        """
        shape = x.size()  # x.shape [b, 2125, 32]
        x = F.softmax(x.reshape(*shape[:-1], 4, self.reg_max + 1), dim=-1)  # [b, 2125, 4, 8]
        x = F.linear(x, self.project.type_as(x)).reshape(*shape[:-1], 4)
        return x
A fixed layer for calculating integral result from distribution.     This layer calculates the target location by:math: 'sum{P(y_i) * y_i}',     P(y_i) denotes the softmax vector that represents the discrete distribution.     y_i denotes the discrete set, usually {0, 1, 2,..., reg_max}     Args:         reg_max (int): the maximal value of the discrete set. Default:16.             you may reset it according to your new dataset or related settings. Forward feature from the regression head to get integral result of bounding box location.         Args:             :param x (Tensor): Features of the regression head, shape (N, 4*(n+1)), n is self.reg_max.             :return x (Tensor): Integral result of box location, i.e., distance offsets from box center in four directions, shape (N, 4). # x.shape [b, 2125, 32] # [b, 2125, 4, 8]