This commit is contained in:
Glenn Jocher 2019-06-15 17:06:58 +02:00
parent 6c77764bba
commit b59532883a
1 changed files with 2 additions and 1 deletions

View File

@ -228,7 +228,7 @@ def train(
pred = model(imgs) pred = model(imgs)
# Compute loss # Compute loss
loss, loss_items = compute_loss(pred, targets, model, giou_loss=False) loss, loss_items = compute_loss(pred, targets, model, giou_loss=opt.giou)
if torch.isnan(loss): if torch.isnan(loss):
print('WARNING: nan loss detected, ending training') print('WARNING: nan loss detected, ending training')
return results return results
@ -325,6 +325,7 @@ if __name__ == '__main__':
parser.add_argument('--backend', default='nccl', type=str, help='distributed backend') parser.add_argument('--backend', default='nccl', type=str, help='distributed backend')
parser.add_argument('--nosave', action='store_true', help='do not save training results') parser.add_argument('--nosave', action='store_true', help='do not save training results')
parser.add_argument('--notest', action='store_true', help='only test final epoch') parser.add_argument('--notest', action='store_true', help='only test final epoch')
parser.add_argument('--giou', action='store_true', help='use GIoU loss instead of xy, wh loss')
parser.add_argument('--evolve', action='store_true', help='run hyperparameter evolution') parser.add_argument('--evolve', action='store_true', help='run hyperparameter evolution')
parser.add_argument('--var', default=0, type=int, help='debug variable') parser.add_argument('--var', default=0, type=int, help='debug variable')
opt = parser.parse_args() opt = parser.parse_args()