Fix load and resume and update autodownload endpoint (#136)
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
parent
6d5123297e
commit
82c849c163
5 changed files with 16 additions and 34 deletions
|
|
@ -17,12 +17,9 @@ from ultralytics.yolo.utils.torch_utils import de_parallel
|
|||
class SegmentationTrainer(v8.detect.DetectionTrainer):
|
||||
|
||||
def load_model(self, model_cfg=None, weights=None, verbose=True):
|
||||
model = SegmentationModel(model_cfg or getattr(weights, 'yaml', None) or weights['model'].yaml,
|
||||
ch=3,
|
||||
nc=self.data["nc"],
|
||||
verbose=verbose)
|
||||
model = SegmentationModel(model_cfg or weights.yaml, ch=3, nc=self.data["nc"], verbose=verbose)
|
||||
if weights:
|
||||
model.load(weights['model'] if isinstance(weights, dict) else weights, verbose)
|
||||
model.load(weights, verbose)
|
||||
return model
|
||||
|
||||
def get_validator(self):
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue