Update .pre-commit-config.yaml (#1026)

ultralytics/yolo/v8/classify/train.py

@@ -16,14 +16,14 @@ class ClassificationTrainer(BaseTrainer):
     def __init__(self, cfg=DEFAULT_CFG, overrides=None):
         if overrides is None:
             overrides = {}
-        overrides["task"] = "classify"
+        overrides['task'] = 'classify'
         super().__init__(cfg, overrides)
 
     def set_model_attributes(self):
-        self.model.names = self.data["names"]
+        self.model.names = self.data['names']
 
     def get_model(self, cfg=None, weights=None, verbose=True):
-        model = ClassificationModel(cfg, nc=self.data["nc"], verbose=verbose and RANK == -1)
+        model = ClassificationModel(cfg, nc=self.data['nc'], verbose=verbose and RANK == -1)
         if weights:
             model.load(weights)
 
@@ -53,11 +53,11 @@ class ClassificationTrainer(BaseTrainer):
 
         model = str(self.model)
         # Load a YOLO model locally, from torchvision, or from Ultralytics assets
-        if model.endswith(".pt"):
+        if model.endswith('.pt'):
             self.model, _ = attempt_load_one_weight(model, device='cpu')
             for p in self.model.parameters():
                 p.requires_grad = True  # for training
-        elif model.endswith(".yaml"):
+        elif model.endswith('.yaml'):
             self.model = self.get_model(cfg=model)
         elif model in torchvision.models.__dict__:
             pretrained = True
@@ -67,15 +67,15 @@ class ClassificationTrainer(BaseTrainer):
 
         return  # dont return ckpt. Classification doesn't support resume
 
-    def get_dataloader(self, dataset_path, batch_size=16, rank=0, mode="train"):
+    def get_dataloader(self, dataset_path, batch_size=16, rank=0, mode='train'):
         loader = build_classification_dataloader(path=dataset_path,
                                                  imgsz=self.args.imgsz,
-                                                 batch_size=batch_size if mode == "train" else (batch_size * 2),
-                                                 augment=mode == "train",
+                                                 batch_size=batch_size if mode == 'train' else (batch_size * 2),
+                                                 augment=mode == 'train',
                                                  rank=rank,
                                                  workers=self.args.workers)
         # Attach inference transforms
-        if mode != "train":
+        if mode != 'train':
             if is_parallel(self.model):
                 self.model.module.transforms = loader.dataset.torch_transforms
             else:
@@ -83,8 +83,8 @@ class ClassificationTrainer(BaseTrainer):
         return loader
 
     def preprocess_batch(self, batch):
-        batch["img"] = batch["img"].to(self.device)
-        batch["cls"] = batch["cls"].to(self.device)
+        batch['img'] = batch['img'].to(self.device)
+        batch['cls'] = batch['cls'].to(self.device)
         return batch
 
     def progress_string(self):
@@ -96,7 +96,7 @@ class ClassificationTrainer(BaseTrainer):
         return v8.classify.ClassificationValidator(self.test_loader, self.save_dir, logger=self.console)
 
     def criterion(self, preds, batch):
-        loss = torch.nn.functional.cross_entropy(preds, batch["cls"], reduction='sum') / self.args.nbs
+        loss = torch.nn.functional.cross_entropy(preds, batch['cls'], reduction='sum') / self.args.nbs
         loss_items = loss.detach()
         return loss, loss_items
 
@@ -112,12 +112,12 @@ class ClassificationTrainer(BaseTrainer):
     #     else:
     #         return keys
 
-    def label_loss_items(self, loss_items=None, prefix="train"):
+    def label_loss_items(self, loss_items=None, prefix='train'):
         """
         Returns a loss dict with labelled training loss items tensor
         """
         # Not needed for classification but necessary for segmentation & detection
-        keys = [f"{prefix}/{x}" for x in self.loss_names]
+        keys = [f'{prefix}/{x}' for x in self.loss_names]
         if loss_items is None:
             return keys
         loss_items = [round(float(loss_items), 5)]
@@ -140,8 +140,8 @@ class ClassificationTrainer(BaseTrainer):
 
 
 def train(cfg=DEFAULT_CFG, use_python=False):
-    model = cfg.model or "yolov8n-cls.pt"  # or "resnet18"
-    data = cfg.data or "mnist160"  # or yolo.ClassificationDataset("mnist")
+    model = cfg.model or 'yolov8n-cls.pt'  # or "resnet18"
+    data = cfg.data or 'mnist160'  # or yolo.ClassificationDataset("mnist")
     device = cfg.device if cfg.device is not None else ''
 
     args = dict(model=model, data=data, device=device)
@@ -153,5 +153,5 @@ def train(cfg=DEFAULT_CFG, use_python=False):
         trainer.train()
 
 
-if __name__ == "__main__":
+if __name__ == '__main__':
     train()