Enable to toggle whether Gem Pooling is trainable or not by requires_grad

Author: yukkyo

mmcls/models/necks/gem.py

@@ -35,10 +35,7 @@ class GeneralizedMeanPooling(nn.Module):
     def __init__(self, p=3., eps=1e-6, clamp=True, p_trainable=True):
         assert p >= 1, "'p' must be a value greater then 1"
         super(GeneralizedMeanPooling, self).__init__()
-        if p_trainable:
-            self.p = Parameter(torch.ones(1) * p)
-        else:
-            self.p = p
+        self.p = Parameter(torch.ones(1) * p, requires_grad=p_trainable)
         self.eps = eps
         self.clamp = clamp
         self.p_trainable = p_trainable

tests/test_models/test_necks.py

@@ -46,7 +46,7 @@ def test_gem_neck():
     neck = GeneralizedMeanPooling()
 
     # default p is trainable
-    assert isinstance(neck.p, torch.nn.Parameter)
+    assert neck.p.requires_grad
 
     # batch_size, num_features, feature_size(2)
     fake_input = torch.rand(1, 16, 24, 24)
@@ -68,8 +68,8 @@ def test_gem_neck():
     # test gem_neck with p_trainable=False
     neck = GeneralizedMeanPooling(p_trainable=False)
 
-    # p is not trainable(float)
-    assert isinstance(neck.p, float)
+    # p is not trainable
+    assert not neck.p.requires_grad
 
     # batch_size, num_features, feature_size(2)
     fake_input = torch.rand(1, 16, 24, 24)