How to convert this code from Pytorch to Tensorflow?

Viewed 18

I am trying to convert code from Pytorch to tensorflow, I don't have much idea about Pytorch, I like to use Tensorflow more, and I need some explanation about this class and some help to write that but in TensorFlow can anyone help me or share useful links Here is Pytorch code :

class DiffGroupNorm(torch.nn.Module):
    def __init__(self, in_channels, groups, lamda=0.01, eps=1e-5, momentum=0.1,
                 affine=True, track_running_stats=True):
        super().__init__()

        self.in_channels = in_channels
        self.groups = groups
        self.lamda = lamda

        self.lin = Linear(in_channels, groups, bias=False)
        self.norm = BatchNorm1d(groups * in_channels, eps, momentum, affine,
                                track_running_stats)
        self.reset_parameters()
    def reset_parameters(self):
            self.lin.reset_parameters()
            self.norm.reset_parameters()
    def forward(self, x: Tensor) -> Tensor:
        """"""
        F, G = self.in_channels, self.groups

        s = self.lin(x).softmax(dim=-1)  # [N, G]
        out = s.unsqueeze(-1) * x.unsqueeze(-2)  # [N, G, F]
        out = self.norm(out.view(-1, G * F)).view(-1, G, F).sum(-2)  # [N, F]

        return x + self.lamda * out
    @staticmethod
    def group_distance_ratio(x: Tensor, y: Tensor, eps: float = 1e-5) -> float:
        num_classes = int(y.max()) + 1

        numerator = 0.
        for i in range(num_classes):
            mask = y == i
            dist = torch.cdist(x[mask].unsqueeze(0), x[~mask].unsqueeze(0))
            numerator += (1 / dist.numel()) * float(dist.sum())
        numerator *= 1 / (num_classes - 1)**2

        denominator = 0.
        for i in range(num_classes):
            mask = y == i
            dist = torch.cdist(x[mask].unsqueeze(0), x[mask].unsqueeze(0))
            denominator += (1 / dist.numel()) * float(dist.sum())
        denominator *= 1 / num_classes

        return numerator / (denominator + eps)


    def __repr__(self) -> str:
        return (f'{self.__class__.__name__}({self.in_channels}, '
                f'groups={self.groups})')
0 Answers
Related