Update train.py (#4136)
* Refactor train.py * Update imports * Update imports * Update optimizer * cleanup
This commit is contained in:
1 parent
264be1a616
commit
63dd65e7ed
3 files changed
+49
-59
No files matched your search
+1
-1
@@ -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))
|
||||
|
||||
|
||||
Reference in new issue
Block a user