This commit is contained in:
Glenn Jocher 2018-11-22 15:04:02 +01:00
parent b9d87be318
commit 154fae4430
1 changed files with 22 additions and 19 deletions

View File

@ -254,11 +254,8 @@ class Darknet(nn.Module):
output.append(x)
layer_outputs.append(x)
self.losses['nT'] /= 3
self.losses['TP'] = 0
self.losses['FP'] = 0
self.losses['FN'] = 0
if is_training and batch_report:
if is_training:
if batch_report:
self.losses['TC'] /= 3 # target category
metrics = torch.zeros(3, len(self.losses['FPe'])) # TP, FP, FN
@ -274,8 +271,14 @@ class Darknet(nn.Module):
self.losses['FP'] = metrics[1].sum()
self.losses['FN'] = metrics[2].sum()
self.losses['metrics'] = metrics
else:
self.losses['TP'] = 0
self.losses['FP'] = 0
self.losses['FN'] = 0
self.losses['nT'] /= 3
self.losses['TC'] = 0
return sum(output) if is_training else torch.cat(output, 1)