Update Time_TravelRephotography/torch_utils/ops/bias_act.py

Time_TravelRephotography/torch_utils/ops/bias_act.py

@@ -88,6 +88,15 @@ def bias_act(x, b=None, dim=1, act='linear', alpha=None, gain=None, clamp=None,
         return _bias_act_cuda(dim=dim, act=act, alpha=alpha, gain=gain, clamp=clamp).apply(x, b)
     return _bias_act_ref(x=x, b=b, dim=dim, act=act, alpha=alpha, gain=gain, clamp=clamp)
 
+class bias_act_relu(nn.Module):
+    def __init__(self, dim, negative_slope=0.2, scale=2 ** 0.5):
+        super().__init__()
+        self.bias = nn.Parameter(torch.zeros(dim))
+        self.negative_slope = negative_slope
+        self.scale = scale
+
+    def forward(self, input):
+        return bias_act(input,  b=self.bias, self.negative_slope, self.scale)
 #----------------------------------------------------------------------------
 
 def _bias_act_ref(x, b=None, dim=1, act='linear', alpha=None, gain=None, clamp=None):
@@ -121,6 +130,7 @@ def _bias_act_ref(x, b=None, dim=1, act='linear', alpha=None, gain=None, clamp=N
         x = x.clamp(-clamp, clamp) # pylint: disable=invalid-unary-operand-type
     return x
 
+
 #----------------------------------------------------------------------------
 
 _bias_act_cuda_cache = dict()