chenxuanhong
48 lines
1.5 KiB
Python
48 lines
1.5 KiB
Python
#!/usr/bin/env python3
|
|
# -*- coding:utf-8 -*-
|
|
#############################################################
|
|
# File: ECA.py
|
|
# Created Date: Tuesday February 23rd 2021
|
|
# Author: Chen Xuanhong
|
|
# Email: chenxuanhongzju@outlook.com
|
|
# Last Modified: Tuesday, 23rd February 2021 9:14:28 pm
|
|
# Modified By: Chen Xuanhong
|
|
# Copyright (c) 2021 Shanghai Jiao Tong University
|
|
#############################################################
|
|
|
|
import math
|
|
import torch
|
|
from torch import nn
|
|
from torch.nn.parameter import Parameter
|
|
|
|
class eca_layer(nn.Module):
|
|
"""Constructs a ECA module.
|
|
Args:
|
|
channel: Number of channels of the input feature map
|
|
k_size: Adaptive selection of kernel size
|
|
"""
|
|
def __init__(self, channel):
|
|
super(eca_layer, self).__init__()
|
|
|
|
b = 1
|
|
gamma = 2
|
|
k_size = int(abs(math.log(channel,2)+b)/gamma)
|
|
k_size = k_size if k_size % 2 else k_size+1
|
|
self.avg_pool = nn.AdaptiveAvgPool2d(1)
|
|
self.conv = nn.Conv1d(1, 1, kernel_size=k_size, padding=(k_size - 1) // 2, bias=False)
|
|
self.sigmoid = nn.Sigmoid()
|
|
|
|
def forward(self, x):
|
|
# x: input features with shape [b, c, h, w]
|
|
# b, c, h, w = x.size()
|
|
|
|
# feature descriptor on the global spatial information
|
|
y = self.avg_pool(x)
|
|
|
|
# Two different branches of ECA module
|
|
y = self.conv(y.squeeze(-1).transpose(-1, -2)).transpose(-1, -2).unsqueeze(-1)
|
|
|
|
# Multi-scale information fusion
|
|
y = self.sigmoid(y)
|
|
|
|
return x * y.expand_as(x) |