From 238a449d0146f96cb69dbd43cfc04248445eae2e Mon Sep 17 00:00:00 2001 From: Erik Gaasedelen Date: Sat, 24 Nov 2018 23:20:17 -0600 Subject: [PATCH] Adaptive pooling for senet Needed to add this line to make it work on other resolution images. --- pretrainedmodels/models/senet.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pretrainedmodels/models/senet.py b/pretrainedmodels/models/senet.py index 20bd122f..bc669b91 100644 --- a/pretrainedmodels/models/senet.py +++ b/pretrainedmodels/models/senet.py @@ -320,7 +320,7 @@ def __init__(self, block, layers, groups, reduction, dropout_p=0.2, downsample_kernel_size=downsample_kernel_size, downsample_padding=downsample_padding ) - self.avg_pool = nn.AvgPool2d(7, stride=1) + self.avg_pool = nn.AdaptiveAvgPool2d(1) self.dropout = nn.Dropout(dropout_p) if dropout_p is not None else None self.last_linear = nn.Linear(512 * block.expansion, num_classes)