Skip to content

Commit 6434dea

Browse files
ekagra-ranjanfmassa
authored andcommitted
Changing to AdaptiveAvgPool2d on Alexnet (#746)
The update allows Alexnet to process images larger or smaller than prescribed imagenet size using adaptive average pooling. Will be useful while finetuning or testing on different resolution images. Similar to #643 and #672. I did not include adaptive avg pool in features or classifier block so that these predefined blocks can be used as it is.
1 parent 9e12551 commit 6434dea

File tree

1 file changed

+2
-0
lines changed

1 file changed

+2
-0
lines changed

torchvision/models/alexnet.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@ def __init__(self, num_classes=1000):
2929
nn.ReLU(inplace=True),
3030
nn.MaxPool2d(kernel_size=3, stride=2),
3131
)
32+
self.avgpool = nn.AdaptiveAvgPool2d((6, 6))
3233
self.classifier = nn.Sequential(
3334
nn.Dropout(),
3435
nn.Linear(256 * 6 * 6, 4096),
@@ -41,6 +42,7 @@ def __init__(self, num_classes=1000):
4142

4243
def forward(self, x):
4344
x = self.features(x)
45+
x = self.avgpool(x)
4446
x = x.view(x.size(0), 256 * 6 * 6)
4547
x = self.classifier(x)
4648
return x

0 commit comments

Comments
 (0)