Source code for mlreco.models.layers.common.nonlinearities

import torch
import torch.nn as nn
import torch.nn.functional as F

# For MinkowskiEngine
import MinkowskiEngine as ME
from MinkowskiNonlinearity import MinkowskiNonlinearityBase

# Custom Nonlinearities

[docs]class MinkowskiLeakyReLU(MinkowskiNonlinearityBase):
    MODULE = nn.LeakyReLU


[docs]    def __repr__(self):
        return self.__class__.__name__ + '(negative_slope = ' + str(self.module.negative_slope) + ')'



[docs]class MinkowskiELU(MinkowskiNonlinearityBase):
    MODULE = nn.ELU


[docs]class MinkowskiMish(nn.Module):
    '''
    Mish Nonlinearity: https://arxiv.org/pdf/1908.08681.pdf
    '''

[docs]    def __init__(self):
        super(MinkowskiMish, self).__init__()



[docs]    def forward(self, input):
        out = F.softplus(input.F)
        out = torch.tanh(out)
        out = out * input.F
        return ME.SparseTensor(
            out,
            coords_key=input.coords_key,
            coords_manager=input.coords_man)



[docs]    def __repr__(self):
        return self.__class__.__name__ + '()'