Fix bugs in TAS: missing ReLU in the end of each searching block

This commit is contained in:
D-X-Y
2020-04-09 06:19:31 +00:00
parent 569b9b406a
commit 076f9c2d41
9 changed files with 108 additions and 13 deletions

View File

@@ -171,7 +171,7 @@ class ResNetBasicblock(nn.Module):
else:
residual, expected_flop_c = inputs, 0
out = additive_func(residual, out_b)
return out, expected_inC_b, sum([expected_flop_a, expected_flop_b, expected_flop_c])
return nn.functional.relu(out, inplace=True), expected_inC_b, sum([expected_flop_a, expected_flop_b, expected_flop_c])
def basic_forward(self, inputs):
basicblock = self.conv_a(inputs)
@@ -243,8 +243,7 @@ class ResNetBottleneck(nn.Module):
else:
residual, expected_flop_c = inputs, 0
out = additive_func(residual, out_1x4)
return out, expected_inC_1x4, sum([expected_flop_1x1, expected_flop_3x3, expected_flop_1x4, expected_flop_c])
return nn.functional.relu(out, inplace=True), expected_inC_1x4, sum([expected_flop_1x1, expected_flop_3x3, expected_flop_1x4, expected_flop_c])
class SearchShapeImagenetResNet(nn.Module):