xclele/yolov10
1
0
Fork 0
mirror of https://github.com/THU-MIG/yolov10.git synced 2025-05-24 14:44:21 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Default classify training to pretrained=True (#3239)

Browse Source
This commit is contained in:
Glenn Jocher 2023-06-18 02:31:05 +02:00 committed by GitHub
parent e78fb683f4
commit 15c90bd404
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 3 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
ultralytics/yolo/cfg/default.yaml
View File

@ -19,7 +19,7 @@ workers: 8 # (int) number of worker threads for data loading (per RANK if DDP)
project: # (str, optional) project name project: # (str, optional) project name
name: # (str, optional) experiment name, results saved to 'project/name' directory name: # (str, optional) experiment name, results saved to 'project/name' directory
exist_ok: False # (bool) whether to overwrite existing experiment exist_ok: False # (bool) whether to overwrite existing experiment
pretrained: False # (bool) whether to use a pretrained model pretrained: True # (bool) whether to use a pretrained model
optimizer: auto # (str) optimizer to use, choices=[SGD, Adam, Adamax, AdamW, NAdam, RAdam, RMSProp, auto] optimizer: auto # (str) optimizer to use, choices=[SGD, Adam, Adamax, AdamW, NAdam, RAdam, RMSProp, auto]
verbose: True # (bool) whether to print verbose output verbose: True # (bool) whether to print verbose output
seed: 0 # (int) random seed for reproducibility seed: 0 # (int) random seed for reproducibility

6
ultralytics/yolo/v8/classify/train.py
View File

@ -33,9 +33,8 @@ class ClassificationTrainer(BaseTrainer):
if weights: if weights:
model.load(weights) model.load(weights)
pretrained = self.args.pretrained
for m in model.modules(): for m in model.modules():
if not pretrained and hasattr(m, 'reset_parameters'): if not self.args.pretrained and hasattr(m, 'reset_parameters'):
m.reset_parameters() m.reset_parameters()
if isinstance(m, torch.nn.Dropout) and self.args.dropout: if isinstance(m, torch.nn.Dropout) and self.args.dropout:
m.p = self.args.dropout # set dropout m.p = self.args.dropout # set dropout
@ -61,8 +60,7 @@ class ClassificationTrainer(BaseTrainer):
elif model.endswith('.yaml'): elif model.endswith('.yaml'):
self.model = self.get_model(cfg=model) self.model = self.get_model(cfg=model)
elif model in torchvision.models.__dict__: elif model in torchvision.models.__dict__:
pretrained = True self.model = torchvision.models.__dict__[model](weights='IMAGENET1K_V1' if self.args.pretrained else None)
self.model = torchvision.models.__dict__[model](weights='IMAGENET1K_V1' if pretrained else None)
else: else:
FileNotFoundError(f'ERROR: model={model} not found locally or online. Please check model name.') FileNotFoundError(f'ERROR: model={model} not found locally or online. Please check model name.')
ClassificationModel.reshape_outputs(self.model, self.data['nc']) ClassificationModel.reshape_outputs(self.model, self.data['nc'])