Update train.py (#4136)

* Refactor train.py

* Update imports

* Update imports

* Update optimizer

* cleanup
This commit is contained in:
Glenn Jocher authored and GitHub committed 2021-07-24 16:11:39 +02:00
1 parent 264be1a616
commit 63dd65e7ed
3 files changed
+49 -59

No files matched your search

+1 -1
View File
@@ -108,7 +108,7 @@ class ComputeLoss:
det = model.module.model[-1] if is_parallel(model) else model.model[-1] # Detect() module
self.balance = {3: [4.0, 1.0, 0.4]}.get(det.nl, [4.0, 1.0, 0.25, 0.06, .02]) # P3-P7
self.ssi = list(det.stride).index(16) if autobalance else 0 # stride 16 index
self.BCEcls, self.BCEobj, self.gr, self.hyp, self.autobalance = BCEcls, BCEobj, model.gr, h, autobalance
self.BCEcls, self.BCEobj, self.gr, self.hyp, self.autobalance = BCEcls, BCEobj, 1.0, h, autobalance
for k in 'na', 'nc', 'nl', 'anchors':
setattr(self, k, getattr(det, k))