Add prec_time to test_net.py and activations count

Summary:
-  This diff introduces changes to run prec_time computation in test_net as well.
- We don't currently have the code to count activations in pycls. This will be useful for users to verify the activations of the upcoming open-sourced models.

Reviewed By: pdollar

Differential Revision: D20971486

fbshipit-source-id: b2319e8fde9e2d862f4fc020bd642b31e498c453

Author: rajprateek

pycls/utils/metrics.py

@@ -72,7 +72,29 @@ def flops_count(model):
             h = (h + 2 * m.padding - m.kernel_size) // m.stride + 1
             w = (w + 2 * m.padding - m.kernel_size) // m.stride + 1
         elif isinstance(m, nn.Linear):
-            count += m.in_features * m.out_features
+            count += m.in_features * m.out_features + m.bias.numel()
+    return count.item()
+
+
+def acts_count(model):
+    """Computes the number of activations statically."""
+    h, w = cfg.TRAIN.IM_SIZE, cfg.TRAIN.IM_SIZE
+    count = 0
+    for n, m in model.named_modules():
+        if isinstance(m, nn.Conv2d):
+            if "se." in n:
+                count += m.out_channels
+                continue
+            h_out = (h + 2 * m.padding[0] - m.kernel_size[0]) // m.stride[0] + 1
+            w_out = (w + 2 * m.padding[1] - m.kernel_size[1]) // m.stride[1] + 1
+            count += np.prod([m.out_channels, h_out, w_out])
+            if ".proj" not in n:
+                h, w = h_out, w_out
+        elif isinstance(m, nn.MaxPool2d):
+            h = (h + 2 * m.padding - m.kernel_size) // m.stride + 1
+            w = (w + 2 * m.padding - m.kernel_size) // m.stride + 1
+        elif isinstance(m, nn.Linear):
+            count += m.out_features
     return count.item()
 
 

tools/test_net.py

@@ -11,13 +11,16 @@
 import sys
 
 import numpy as np
+import pycls.core.losses as losses
 import pycls.core.model_builder as model_builder
 import pycls.datasets.loader as loader
+import pycls.utils.benchmark as bu
 import pycls.utils.checkpoint as cu
 import pycls.utils.distributed as du
 import pycls.utils.logging as lu
 import pycls.utils.metrics as mu
 import pycls.utils.multiprocessing as mpu
+import pycls.utils.net as nu
 import torch
 from pycls.core.config import assert_and_infer_cfg, cfg
 from pycls.utils.meters import TestMeter
@@ -49,6 +52,7 @@ def log_model_info(model):
     logger.info("Model:\n{}".format(model))
     logger.info("Params: {:,}".format(mu.params_count(model)))
     logger.info("Flops: {:,}".format(mu.flops_count(model)))
+    logger.info("Acts: {:,}".format(mu.acts_count(model)))
 
 
 @torch.no_grad()
@@ -89,6 +93,13 @@ def test_model():
     model = model_builder.build_model()
     log_model_info(model)
 
+    # Compute precise time
+    if cfg.PREC_TIME.ENABLED:
+        logger.info("Computing precise time...")
+        loss_fun = losses.get_loss_fun()
+        bu.compute_precise_time(model, loss_fun)
+        nu.reset_bn_stats(model)
+
     # Load model weights
     cu.load_checkpoint(cfg.TEST.WEIGHTS, model)
     logger.info("Loaded model weights from: {}".format(cfg.TEST.WEIGHTS))

tools/train_net.py

@@ -61,6 +61,7 @@ def log_model_info(model):
     logger.info("Model:\n{}".format(model))
     logger.info("Params: {:,}".format(mu.params_count(model)))
     logger.info("Flops: {:,}".format(mu.flops_count(model)))
+    logger.info("Acts: {:,}".format(mu.acts_count(model)))
 
 
 def train_epoch(train_loader, model, loss_fun, optimizer, train_meter, cur_epoch):