|
|
|
@ -38,6 +38,9 @@ from torch import nn
|
|
|
|
|
import torch.nn.functional as F
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
from .create_act import create_act_layer
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class EcaModule(nn.Module):
|
|
|
|
|
"""Constructs an ECA module.
|
|
|
|
|
|
|
|
|
@ -48,20 +51,27 @@ class EcaModule(nn.Module):
|
|
|
|
|
refer to original paper https://arxiv.org/pdf/1910.03151.pdf
|
|
|
|
|
(default=None. if channel size not given, use k_size given for kernel size.)
|
|
|
|
|
kernel_size: Adaptive selection of kernel size (default=3)
|
|
|
|
|
gamm: used in kernel_size calc, see above
|
|
|
|
|
beta: used in kernel_size calc, see above
|
|
|
|
|
act_layer: optional non-linearity after conv, enables conv bias, this is an experiment
|
|
|
|
|
gate_layer: gating non-linearity to use
|
|
|
|
|
"""
|
|
|
|
|
def __init__(self, channels=None, kernel_size=3, gamma=2, beta=1):
|
|
|
|
|
def __init__(self, channels=None, kernel_size=3, gamma=2, beta=1, act_layer=None, gate_layer='sigmoid'):
|
|
|
|
|
super(EcaModule, self).__init__()
|
|
|
|
|
assert kernel_size % 2 == 1
|
|
|
|
|
if channels is not None:
|
|
|
|
|
t = int(abs(math.log(channels, 2) + beta) / gamma)
|
|
|
|
|
kernel_size = max(t if t % 2 else t + 1, 3)
|
|
|
|
|
|
|
|
|
|
self.conv = nn.Conv1d(1, 1, kernel_size=kernel_size, padding=(kernel_size - 1) // 2, bias=False)
|
|
|
|
|
assert kernel_size % 2 == 1
|
|
|
|
|
has_act = act_layer is not None
|
|
|
|
|
self.conv = nn.Conv1d(1, 1, kernel_size=kernel_size, padding=(kernel_size - 1) // 2, bias=has_act)
|
|
|
|
|
self.act = create_act_layer(act_layer) if has_act else nn.Identity()
|
|
|
|
|
self.gate = create_act_layer(gate_layer)
|
|
|
|
|
|
|
|
|
|
def forward(self, x):
|
|
|
|
|
y = x.mean((2, 3)).view(x.shape[0], 1, -1) # view for 1d conv
|
|
|
|
|
y = self.conv(y)
|
|
|
|
|
y = y.view(x.shape[0], -1, 1, 1).sigmoid()
|
|
|
|
|
y = self.act(y) # NOTE: usually a no-op, added for experimentation
|
|
|
|
|
y = self.gate(y).view(x.shape[0], -1, 1, 1)
|
|
|
|
|
return x * y.expand_as(x)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
@ -86,27 +96,35 @@ class CecaModule(nn.Module):
|
|
|
|
|
refer to original paper https://arxiv.org/pdf/1910.03151.pdf
|
|
|
|
|
(default=None. if channel size not given, use k_size given for kernel size.)
|
|
|
|
|
kernel_size: Adaptive selection of kernel size (default=3)
|
|
|
|
|
gamm: used in kernel_size calc, see above
|
|
|
|
|
beta: used in kernel_size calc, see above
|
|
|
|
|
act_layer: optional non-linearity after conv, enables conv bias, this is an experiment
|
|
|
|
|
gate_layer: gating non-linearity to use
|
|
|
|
|
"""
|
|
|
|
|
|
|
|
|
|
def __init__(self, channels=None, kernel_size=3, gamma=2, beta=1):
|
|
|
|
|
def __init__(self, channels=None, kernel_size=3, gamma=2, beta=1, act_layer=None, gate_layer='sigmoid'):
|
|
|
|
|
super(CecaModule, self).__init__()
|
|
|
|
|
assert kernel_size % 2 == 1
|
|
|
|
|
if channels is not None:
|
|
|
|
|
t = int(abs(math.log(channels, 2) + beta) / gamma)
|
|
|
|
|
kernel_size = max(t if t % 2 else t + 1, 3)
|
|
|
|
|
has_act = act_layer is not None
|
|
|
|
|
assert kernel_size % 2 == 1
|
|
|
|
|
|
|
|
|
|
# PyTorch circular padding mode is buggy as of pytorch 1.4
|
|
|
|
|
# see https://github.com/pytorch/pytorch/pull/17240
|
|
|
|
|
# implement manual circular padding
|
|
|
|
|
self.conv = nn.Conv1d(1, 1, kernel_size=kernel_size, padding=0, bias=False)
|
|
|
|
|
self.padding = (kernel_size - 1) // 2
|
|
|
|
|
self.conv = nn.Conv1d(1, 1, kernel_size=kernel_size, padding=0, bias=has_act)
|
|
|
|
|
self.act = create_act_layer(act_layer) if has_act else nn.Identity()
|
|
|
|
|
self.gate = create_act_layer(gate_layer)
|
|
|
|
|
|
|
|
|
|
def forward(self, x):
|
|
|
|
|
y = x.mean((2, 3)).view(x.shape[0], 1, -1)
|
|
|
|
|
# Manually implement circular padding, F.pad does not seemed to be bugged
|
|
|
|
|
y = F.pad(y, (self.padding, self.padding), mode='circular')
|
|
|
|
|
y = self.conv(y)
|
|
|
|
|
y = y.view(x.shape[0], -1, 1, 1).sigmoid()
|
|
|
|
|
y = self.act(y) # NOTE: usually a no-op, added for experimentation
|
|
|
|
|
y = self.gate(y).view(x.shape[0], -1, 1, 1)
|
|
|
|
|
return x * y.expand_as(x)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|