Commit d0ebcec4 authored by Myle Ott's avatar Myle Ott Committed by Facebook Github Bot
Browse files

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
parent 38f1dee9

train.py

+5 −1
Original line number Diff line number Diff line
@@ -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