Fix incorrect epoch in the checkpoint (#159)

eldarkurtic · web-flow · commit cf441a7dacd2 · 2021-04-07T21:34:58.000-04:00
The checkpoint created at the end of training is saved with the wrong epoch number because of the `epoch += 1` in the training loop.
diff --git a/scripts/pytorch_vision.py b/scripts/pytorch_vision.py
@@ -1193,7 +1193,7 @@ def train(args, model, train_loader, val_loader, input_shape, save_dir, loggers)
             # only convert qat -> quantized ONNX graph for finalized model
             # TODO: change this to all checkpoints when conversion times improve
             _save_model_training(
-                model, optim, input_shape, "model", save_dir, epoch, val_res, True
+                model, optim, input_shape, "model", save_dir, epoch - 1, val_res, True
             )
 
             LOGGER.info("layer sparsities:")

Original file line number	Diff line number	Diff line change
`@@ -1193,7 +1193,7 @@ def train(args, model, train_loader, val_loader, input_shape, save_dir, loggers)`
`1193`	`1193`	`# only convert qat -> quantized ONNX graph for finalized model`
`1194`	`1194`	`# TODO: change this to all checkpoints when conversion times improve`
`1195`	`1195`	`_save_model_training(`
`1196`		`- model, optim, input_shape, "model", save_dir, epoch, val_res, True`
	`1196`	`+ model, optim, input_shape, "model", save_dir, epoch - 1, val_res, True`
`1197`	`1197`	`)`
`1198`	`1198`
`1199`	`1199`	`LOGGER.info("layer sparsities:")`