Print model and number of trained params

Summary: Pull Request resolved: https://github.com/pytorch/fairseq/pull/469

Differential Revision: D13802945

Pulled By: myleott

fbshipit-source-id: b6976506a8336b96ee40505c4a7638541cc99c95

train.py

@@ -44,8 +44,12 @@ def main(args):
     # Build model and criterion
     model = task.build_model(args)
     criterion = task.build_criterion(args)
+    print(model)
     print('| model {}, criterion {}'.format(args.arch, criterion.__class__.__name__))
-    print('| num. model params: {}'.format(sum(p.numel() for p in model.parameters())))
+    print('| num. model params: {} (num. trained: {})'.format(
+        sum(p.numel() for p in model.parameters()),
+        sum(p.numel() for p in model.parameters() if p.requires_grad),
+    ))
 
     # Make a dummy batch to (i) warm the caching allocator and (ii) as a
     # placeholder DistributedDataParallel when there's an uneven number of