Add support for profiling EfficientNet (#483)

Summary:
Pull Request resolved: #483

- Add `flops` and `activations` to `Conv2dSamePadding`
- The EfficientNet forward squeezed the final output - this works fine for batches with > 1 elements, but for unit batches, it removes the batch dimension as well. This resulted in incorrect profiler behavior and has been fixed.

Reviewed By: vreis

Differential Revision: D21094291

fbshipit-source-id: e76a7c5398c521c11e00bca4d428138430bc7344

classy_vision/models/efficientnet.py

@@ -143,6 +143,8 @@ def __init__(
         pad_w = max(
             (out_w - 1) * self.stride[1] + (kernel_w - 1) * dilation_w + 1 - image_w, 0
         )
+        self.out_h = out_h
+        self.out_w = out_w
         self.same_padding = None
         if pad_h > 0 or pad_w > 0:
             self.same_padding = nn.ZeroPad2d(
@@ -168,6 +170,22 @@ def forward(self, x):
         )
         return x
 
+    def flops(self, x):
+        batchsize_per_replica = x.size()[0]
+        return (
+            batchsize_per_replica
+            * self.in_channels
+            * self.out_channels
+            * self.kernel_size[0]
+            * self.kernel_size[1]
+            * self.out_h
+            * self.out_w
+            / self.groups
+        )
+
+    def activations(self, x, out):
+        return out.numel()
+
 
 class MBConvBlock(nn.Module):
     """
@@ -527,7 +545,7 @@ def forward(self, inputs):
         outputs = self.trunk_output(outputs)
 
         # Average Pooling
-        outputs = self.avg_pooling(outputs).squeeze()
+        outputs = self.avg_pooling(outputs).view(outputs.size(0), -1)
 
         # Dropout
         if self.dropout is not None: