mirrors/pytorch-image-models
1
0
Fork 0
mirror of https://github.com/huggingface/pytorch-image-models.git synced 2025-06-03 15:01:08 +08:00
Code Issues Projects Releases Wiki Activity

Change block avgpool in senets to mean for performance issues with NVIDIA and AMP especially

Browse Source
This commit is contained in:
Ross Wightman 2019-04-05 10:53:13 -07:00
parent 5180f94c7e
commit 58571e992e
1 changed files with 3 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
models/senet.py
View File

@ -38,7 +38,7 @@ class SEModule(nn.Module):
def __init__(self, channels, reduction):
super(SEModule, self).__init__()
self.avg_pool = nn.AdaptiveAvgPool2d(1)
#self.avg_pool = nn.AdaptiveAvgPool2d(1)
self.fc1 = nn.Conv2d(
channels, channels // reduction, kernel_size=1, padding=0)
self.relu = nn.ReLU(inplace=True)
@ -48,7 +48,8 @@ class SEModule(nn.Module):
def forward(self, x):
module_input = x
x = self.avg_pool(x)
#x = self.avg_pool(x)
x = x.view(x.size(0), x.size(1), -1).mean(-1).view(x.size(0), x.size(1), 1, 1)
x = self.fc1(x)
x = self.relu(x)
x = self.fc2(x)