Move loss calculation to head (#2874)
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Laughing-q <1185102784@qq.com> Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com> Co-authored-by: Laughing <61612323+Laughing-q@users.noreply.github.com>
This commit is contained in:
parent
7f077f7654
commit
facb7861cf
9 changed files with 417 additions and 348 deletions
|
|
@ -157,12 +157,12 @@ class BaseValidator:
|
|||
|
||||
# Inference
|
||||
with dt[1]:
|
||||
preds = model(batch['img'])
|
||||
preds = model(batch['img'], augment=self.args.augment)
|
||||
|
||||
# Loss
|
||||
with dt[2]:
|
||||
if self.training:
|
||||
self.loss += trainer.criterion(preds, batch)[1]
|
||||
self.loss += model.loss(batch, preds)[1]
|
||||
|
||||
# Postprocess
|
||||
with dt[3]:
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue