PVD/modules/se.py

24 lines
657 B
Python
Raw Normal View History

2021-10-19 20:54:46 +00:00
import torch
2023-04-11 09:12:58 +00:00
import torch.nn as nn
__all__ = ["SE3d"]
2021-10-19 20:54:46 +00:00
class Swish(nn.Module):
2023-04-11 09:12:58 +00:00
def forward(self, x):
return x * torch.sigmoid(x)
2021-10-19 20:54:46 +00:00
class SE3d(nn.Module):
def __init__(self, channel, reduction=8, use_relu=False):
super().__init__()
self.fc = nn.Sequential(
nn.Linear(channel, channel // reduction, bias=False),
2023-04-11 09:12:58 +00:00
nn.ReLU(True) if use_relu else Swish(),
2021-10-19 20:54:46 +00:00
nn.Linear(channel // reduction, channel, bias=False),
2023-04-11 09:12:58 +00:00
nn.Sigmoid(),
2021-10-19 20:54:46 +00:00
)
def forward(self, inputs):
return inputs * self.fc(inputs.mean(-1).mean(-1).mean(-1)).view(inputs.shape[0], inputs.shape[1], 1, 1, 1)