Allow averaging update based checkpoints given upper bound
Summary: This script currently supports averaging last-n updates, last-n epochs and n epochs before an upper-bound. While the code works for n updates before an upper bound, it was not exposed. This diff changes that.
@@ -106,8 +106,11 @@ def main():
help='if set, will try to find checkpoints with names checkpoint_ee_xx.pt in the path specified by input, '
'and average last this many of them.')
parser.add_argument('--checkpoint-upper-bound', type=int,
- help='when using --num-epoch-checkpoints, this will set an upper bound on which checkpoint to use, '
- 'e.g., with --num-epoch-checkpoints=10 --checkpoint-upper-bound=50, checkpoints 41-50 would be averaged.')
+ help='when using --num-epoch-checkpoints, this will set an upper bound on which epoch to use, '
+ 'when using --num-update-checkpoints, this will set an upper bound on which update to use'
+ 'e.g., with --num-epoch-checkpoints=10 --checkpoint-upper-bound=50, checkpoints 41-50 would be averaged.'
+ 'e.g., with --num-update-checkpoints=10 --checkpoint-upper-bound=50000, checkpoints 40500-50000 would be averaged assuming --save-interval-updates 500'
+ )
# fmt: on
args = parser.parse_args()
print(args)
@@ -120,8 +123,8 @@ def main():
elif args.num_epoch_checkpoints is not None:
num = args.num_epoch_checkpoints
- assert args.checkpoint_upper_bound is None or args.num_epoch_checkpoints is not None, \
- '--checkpoint-upper-bound requires --num-epoch-checkpoints'
+ assert args.checkpoint_upper_bound is None or (args.num_epoch_checkpoints is not None or args.num_update_checkpoints is not None), \
+ '--checkpoint-upper-bound requires --num-epoch-checkpoints or --num-update-checkpoints'
assert args.num_epoch_checkpoints is None or args.num_update_checkpoints is None, \
'Cannot combine --num-epoch-checkpoints and --num-update-checkpoints'
