activations.py 3.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103
  1. # YOLOv5 🚀 by Ultralytics, GPL-3.0 license
  2. """
  3. Activation functions
  4. """
  5. import torch
  6. import torch.nn as nn
  7. import torch.nn.functional as F
  8. class SiLU(nn.Module):
  9. # SiLU activation https://arxiv.org/pdf/1606.08415.pdf
  10. @staticmethod
  11. def forward(x):
  12. return x * torch.sigmoid(x)
  13. class Hardswish(nn.Module):
  14. # Hard-SiLU activation
  15. @staticmethod
  16. def forward(x):
  17. # return x * F.hardsigmoid(x) # for TorchScript and CoreML
  18. return x * F.hardtanh(x + 3, 0.0, 6.0) / 6.0 # for TorchScript, CoreML and ONNX
  19. class Mish(nn.Module):
  20. # Mish activation https://github.com/digantamisra98/Mish
  21. @staticmethod
  22. def forward(x):
  23. return x * F.softplus(x).tanh()
  24. class MemoryEfficientMish(nn.Module):
  25. # Mish activation memory-efficient
  26. class F(torch.autograd.Function):
  27. @staticmethod
  28. def forward(ctx, x):
  29. ctx.save_for_backward(x)
  30. return x.mul(torch.tanh(F.softplus(x))) # x * tanh(ln(1 + exp(x)))
  31. @staticmethod
  32. def backward(ctx, grad_output):
  33. x = ctx.saved_tensors[0]
  34. sx = torch.sigmoid(x)
  35. fx = F.softplus(x).tanh()
  36. return grad_output * (fx + x * sx * (1 - fx * fx))
  37. def forward(self, x):
  38. return self.F.apply(x)
  39. class FReLU(nn.Module):
  40. # FReLU activation https://arxiv.org/abs/2007.11824
  41. def __init__(self, c1, k=3): # ch_in, kernel
  42. super().__init__()
  43. self.conv = nn.Conv2d(c1, c1, k, 1, 1, groups=c1, bias=False)
  44. self.bn = nn.BatchNorm2d(c1)
  45. def forward(self, x):
  46. return torch.max(x, self.bn(self.conv(x)))
  47. class AconC(nn.Module):
  48. r""" ACON activation (activate or not)
  49. AconC: (p1*x-p2*x) * sigmoid(beta*(p1*x-p2*x)) + p2*x, beta is a learnable parameter
  50. according to "Activate or Not: Learning Customized Activation" <https://arxiv.org/pdf/2009.04759.pdf>.
  51. """
  52. def __init__(self, c1):
  53. super().__init__()
  54. self.p1 = nn.Parameter(torch.randn(1, c1, 1, 1))
  55. self.p2 = nn.Parameter(torch.randn(1, c1, 1, 1))
  56. self.beta = nn.Parameter(torch.ones(1, c1, 1, 1))
  57. def forward(self, x):
  58. dpx = (self.p1 - self.p2) * x
  59. return dpx * torch.sigmoid(self.beta * dpx) + self.p2 * x
  60. class MetaAconC(nn.Module):
  61. r""" ACON activation (activate or not)
  62. MetaAconC: (p1*x-p2*x) * sigmoid(beta*(p1*x-p2*x)) + p2*x, beta is generated by a small network
  63. according to "Activate or Not: Learning Customized Activation" <https://arxiv.org/pdf/2009.04759.pdf>.
  64. """
  65. def __init__(self, c1, k=1, s=1, r=16): # ch_in, kernel, stride, r
  66. super().__init__()
  67. c2 = max(r, c1 // r)
  68. self.p1 = nn.Parameter(torch.randn(1, c1, 1, 1))
  69. self.p2 = nn.Parameter(torch.randn(1, c1, 1, 1))
  70. self.fc1 = nn.Conv2d(c1, c2, k, s, bias=True)
  71. self.fc2 = nn.Conv2d(c2, c1, k, s, bias=True)
  72. # self.bn1 = nn.BatchNorm2d(c2)
  73. # self.bn2 = nn.BatchNorm2d(c1)
  74. def forward(self, x):
  75. y = x.mean(dim=2, keepdims=True).mean(dim=3, keepdims=True)
  76. # batch-size 1 bug/instabilities https://github.com/ultralytics/yolov5/issues/2891
  77. # beta = torch.sigmoid(self.bn2(self.fc2(self.bn1(self.fc1(y))))) # bug/unstable
  78. beta = torch.sigmoid(self.fc2(self.fc1(y))) # bug patch BN layers removed
  79. dpx = (self.p1 - self.p2) * x
  80. return dpx * torch.sigmoid(beta * dpx) + self.p2 * x