fix bug (match channel)

bearpaw · bearpaw · commit f7fe8d67bfbc · 2017-05-15T10:48:32.000+08:00
diff --git a/models/modules.py b/models/modules.py
@@ -83,14 +83,14 @@ def forward(self, x):
 
         return out
 
-# Softmax Attention 
-model = SoftmaxAttention(3).cuda()
-x = torch.randn(2,3,32,32)
-out = model(Variable(x.cuda()))
-print(out.size())
-
-# Sigmoid Attention
-model = SigmoidAttention(3).cuda()
-x = torch.randn(2,3,32,32)
-out = model(Variable(x.cuda()))
-print(out.size())
+# # Softmax Attention 
+# model = SoftmaxAttention(3).cuda()
+# x = torch.randn(2,3,32,32)
+# out = model(Variable(x.cuda()))
+# print(out.size())
+
+# # Sigmoid Attention
+# model = SigmoidAttention(3).cuda()
+# x = torch.randn(2,3,32,32)
+# out = model(Variable(x.cuda()))
+# print(out.size())
diff --git a/models/ressoftattnet.py b/models/ressoftattnet.py
@@ -26,8 +26,8 @@ def __init__(self, block, layers, num_classes=1000):
         self.layer1 = self._make_layer(block, 16, layers[0])
         self.layer2 = self._make_layer(block, 32, layers[1], stride=2)
         self.layer3 = self._make_layer(block, 64, layers[2], stride=2)
-        self.att = SoftmaxAttention(64)
-        self.bn = nn.BatchNorm2d(64)
+        self.att = SoftmaxAttention(64 * block.expansion)
+        self.bn = nn.BatchNorm2d(64 * block.expansion)
         self.avgpool = nn.AvgPool2d(8)
         self.fc = nn.Linear(64 * block.expansion, num_classes)
 
