Multiple Enhancement to Make MMENGINE walk further

.pre-commit-config.yaml

@@ -12,10 +12,14 @@ repos:
     rev: 5.11.5
     hooks:
       - id: isort
-  - repo: https://github.com/pre-commit/mirrors-yapf
-    rev: v0.32.0
+  - repo: local
     hooks:
       - id: yapf
+        name: yapf
+        entry: yapf
+        language: system
+        types: [python]
+        args: ["-i"]
   - repo: https://github.com/pre-commit/pre-commit-hooks
     rev: v5.0.0
     hooks:

examples/distributed_training.py

@@ -42,11 +42,10 @@ def compute_metrics(self, results):
 
 def parse_args():
     parser = argparse.ArgumentParser(description='Distributed Training')
-    parser.add_argument(
-        '--launcher',
-        choices=['none', 'pytorch', 'slurm', 'mpi'],
-        default='none',
-        help='job launcher')
+    parser.add_argument('--launcher',
+                        choices=['none', 'pytorch', 'slurm', 'mpi'],
+                        default='none',
+                        help='job launcher')
     parser.add_argument('--local_rank', type=int, default=0)
 
     args = parser.parse_args()
@@ -73,16 +72,14 @@ def main():
         transform=transforms.Compose(
             [transforms.ToTensor(),
              transforms.Normalize(**norm_cfg)]))
-    train_dataloader = dict(
-        batch_size=32,
-        dataset=train_set,
-        sampler=dict(type='DefaultSampler', shuffle=True),
-        collate_fn=dict(type='default_collate'))
-    val_dataloader = dict(
-        batch_size=32,
-        dataset=valid_set,
-        sampler=dict(type='DefaultSampler', shuffle=False),
-        collate_fn=dict(type='default_collate'))
+    train_dataloader = dict(batch_size=32,
+                            dataset=train_set,
+                            sampler=dict(type='DefaultSampler', shuffle=True),
+                            collate_fn=dict(type='default_collate'))
+    val_dataloader = dict(batch_size=32,
+                          dataset=valid_set,
+                          sampler=dict(type='DefaultSampler', shuffle=False),
+                          collate_fn=dict(type='default_collate'))
     runner = Runner(
         model=MMResNet50(),
         work_dir='./work_dirs',

examples/distributed_training_with_flexible_runner.py

@@ -70,16 +70,14 @@ def main():
         transform=transforms.Compose(
             [transforms.ToTensor(),
              transforms.Normalize(**norm_cfg)]))
-    train_dataloader = dict(
-        batch_size=128,
-        dataset=train_set,
-        sampler=dict(type='DefaultSampler', shuffle=True),
-        collate_fn=dict(type='default_collate'))
-    val_dataloader = dict(
-        batch_size=128,
-        dataset=valid_set,
-        sampler=dict(type='DefaultSampler', shuffle=False),
-        collate_fn=dict(type='default_collate'))
+    train_dataloader = dict(batch_size=128,
+                            dataset=train_set,
+                            sampler=dict(type='DefaultSampler', shuffle=True),
+                            collate_fn=dict(type='default_collate'))
+    val_dataloader = dict(batch_size=128,
+                          dataset=valid_set,
+                          sampler=dict(type='DefaultSampler', shuffle=False),
+                          collate_fn=dict(type='default_collate'))
 
     if args.use_deepspeed:
         strategy = dict(
@@ -97,30 +95,28 @@ def main():
             # bf16=dict(
             #     enabled=True,
             # ),
-            zero_optimization=dict(
-                stage=3,
-                allgather_partitions=True,
-                reduce_scatter=True,
-                allgather_bucket_size=50000000,
-                reduce_bucket_size=50000000,
-                overlap_comm=True,
-                contiguous_gradients=True,
-                cpu_offload=False),
+            zero_optimization=dict(stage=3,
+                                   allgather_partitions=True,
+                                   reduce_scatter=True,
+                                   allgather_bucket_size=50000000,
+                                   reduce_bucket_size=50000000,
+                                   overlap_comm=True,
+                                   contiguous_gradients=True,
+                                   cpu_offload=False),
         )
-        optim_wrapper = dict(
-            type='DeepSpeedOptimWrapper',
-            optimizer=dict(type='AdamW', lr=1e-3))
+        optim_wrapper = dict(type='DeepSpeedOptimWrapper',
+                             optimizer=dict(type='AdamW', lr=1e-3))
     elif args.use_fsdp:
         from functools import partial
 
         from torch.distributed.fsdp.wrap import size_based_auto_wrap_policy
-        size_based_auto_wrap_policy = partial(
-            size_based_auto_wrap_policy, min_num_params=1e7)
+        size_based_auto_wrap_policy = partial(size_based_auto_wrap_policy,
+                                              min_num_params=1e7)
         strategy = dict(
             type='FSDPStrategy',
             model_wrapper=dict(auto_wrap_policy=size_based_auto_wrap_policy))
-        optim_wrapper = dict(
-            type='AmpOptimWrapper', optimizer=dict(type='AdamW', lr=1e-3))
+        optim_wrapper = dict(type='AmpOptimWrapper',
+                             optimizer=dict(type='AdamW', lr=1e-3))
     elif args.use_colossalai:
         from colossalai.tensor.op_wrapper import colo_op_impl
 
@@ -142,20 +138,21 @@ def main():
         optim_wrapper = dict(optimizer=dict(type='HybridAdam', lr=1e-3))
     else:
         strategy = None
-        optim_wrapper = dict(
-            type='AmpOptimWrapper', optimizer=dict(type='AdamW', lr=1e-3))
-
-    runner = FlexibleRunner(
-        model=MMResNet50(),
-        work_dir='./work_dirs',
-        strategy=strategy,
-        train_dataloader=train_dataloader,
-        optim_wrapper=optim_wrapper,
-        param_scheduler=dict(type='LinearLR'),
-        train_cfg=dict(by_epoch=True, max_epochs=10, val_interval=1),
-        val_dataloader=val_dataloader,
-        val_cfg=dict(),
-        val_evaluator=dict(type=Accuracy))
+        optim_wrapper = dict(type='AmpOptimWrapper',
+                             optimizer=dict(type='AdamW', lr=1e-3))
+
+    runner = FlexibleRunner(model=MMResNet50(),
+                            work_dir='./work_dirs',
+                            strategy=strategy,
+                            train_dataloader=train_dataloader,
+                            optim_wrapper=optim_wrapper,
+                            param_scheduler=dict(type='LinearLR'),
+                            train_cfg=dict(by_epoch=True,
+                                           max_epochs=10,
+                                           val_interval=1),
+                            val_dataloader=val_dataloader,
+                            val_cfg=dict(),
+                            val_evaluator=dict(type=Accuracy))
     runner.train()
 
 

examples/llama2/fsdp_finetune.py

@@ -92,17 +92,14 @@ def parse_args():
 def train():
     args = parse_args()
     # Setup distributed related component in Strategy.
-    strategy = FSDPStrategy(
-        model_wrapper=dict(
-            auto_wrap_policy=partial(
-                transformer_auto_wrap_policy,
-                transformer_layer_cls={LlamaDecoderLayer})),
-        state_dict_cfg='full',
-        env_kwargs=dict(randomness=dict(seed=42)))
-    visualizer = Visualizer(
-        name='mmengine',
-        save_dir=args.output_dir,
-        vis_backends=[dict(type=WandbVisBackend)])
+    strategy = FSDPStrategy(model_wrapper=dict(
+        auto_wrap_policy=partial(transformer_auto_wrap_policy,
+                                 transformer_layer_cls={LlamaDecoderLayer})),
+                            state_dict_cfg='full',
+                            env_kwargs=dict(randomness=dict(seed=42)))
+    visualizer = Visualizer(name='mmengine',
+                            save_dir=args.output_dir,
+                            vis_backends=[dict(type=WandbVisBackend)])
 
     # Prepare model
     tokenizer = LlamaTokenizer.from_pretrained(args.checkpoint)
@@ -112,21 +109,20 @@ def train():
     model.train()
 
     # Prepare dataset
-    train_dataset = AlpacaDataset(
-        tokenizer=tokenizer, data_path=args.data_root)
-    train_dataloader = DataLoader(
-        train_dataset,
-        batch_size=args.batch_size,
-        sampler=DefaultSampler(train_dataset, seed=0),
-        collate_fn=default_data_collator,
-        drop_last=True)
+    train_dataset = AlpacaDataset(tokenizer=tokenizer,
+                                  data_path=args.data_root)
+    train_dataloader = DataLoader(train_dataset,
+                                  batch_size=args.batch_size,
+                                  sampler=DefaultSampler(train_dataset,
+                                                         seed=0),
+                                  collate_fn=default_data_collator,
+                                  drop_last=True)
 
     # Get the prepared model, scheduler and optimizer from strategy
     epoch_length = len(train_dataloader)
     max_iters = epoch_length * args.max_epoch
-    optim_cfg = dict(
-        optimizer=dict(type=AdamW, lr=1e-4, weight_decay=0.0),
-        accumulative_counts=ORI_BATCH_SIZE / args.batch_size)
+    optim_cfg = dict(optimizer=dict(type=AdamW, lr=1e-4, weight_decay=0.0),
+                     accumulative_counts=ORI_BATCH_SIZE / args.batch_size)
     scheduler_cfgs = [dict(type=StepLR, step_size=1, gamma=0.85)]
     model, optimizer, schedulers = strategy.prepare(
         model,

examples/llama2/generate.py

@@ -30,7 +30,6 @@ def parse_args():
     with torch.no_grad():
         generate_ids = model.generate(inputs.input_ids.cuda(), max_length=300)
     print(
-        tokenizer.batch_decode(
-            generate_ids,
-            skip_special_tokens=True,
-            clean_up_tokenization_spaces=False)[0])
+        tokenizer.batch_decode(generate_ids,
+                               skip_special_tokens=True,
+                               clean_up_tokenization_spaces=False)[0])

examples/segmentation/train.py

@@ -40,8 +40,9 @@ def __init__(self,
                  mask_folder,
                  transform=None,
                  target_transform=None):
-        super().__init__(
-            root, transform=transform, target_transform=target_transform)
+        super().__init__(root,
+                         transform=transform,
+                         target_transform=target_transform)
         self.img_folder = img_folder
         self.mask_folder = mask_folder
         self.images = list(
@@ -72,8 +73,9 @@ def __getitem__(self, index):
 
         if self.target_transform is not None:
             labels = self.target_transform(labels)
-        data_samples = dict(
-            labels=labels, img_path=img_path, mask_path=mask_path)
+        data_samples = dict(labels=labels,
+                            img_path=img_path,
+                            mask_path=mask_path)
         return img, data_samples
 
     def __len__(self):
@@ -102,8 +104,8 @@ def process(self, data_batch, data_samples):
         intersect = (labels == preds).sum()
         union = (torch.logical_or(preds, labels)).sum()
         iou = (intersect / union).cpu()
-        self.results.append(
-            dict(batch_size=len(labels), iou=iou * len(labels)))
+        self.results.append(dict(batch_size=len(labels),
+                                 iou=iou * len(labels)))
 
     def compute_metrics(self, results):
         total_iou = sum(result['iou'] for result in self.results)
@@ -151,18 +153,16 @@ def after_val_iter(self,
                             osp.join(saved_dir, osp.basename(img_path)))
             shutil.copyfile(mask_path,
                             osp.join(saved_dir, osp.basename(mask_path)))
-            cv2.imwrite(
-                osp.join(saved_dir, f'pred_{osp.basename(img_path)}'),
-                pred_mask)
+            cv2.imwrite(osp.join(saved_dir, f'pred_{osp.basename(img_path)}'),
+                        pred_mask)
 
 
 def parse_args():
     parser = argparse.ArgumentParser(description='Distributed Training')
-    parser.add_argument(
-        '--launcher',
-        choices=['none', 'pytorch', 'slurm', 'mpi'],
-        default='none',
-        help='job launcher')
+    parser.add_argument('--launcher',
+                        choices=['none', 'pytorch', 'slurm', 'mpi'],
+                        default='none',
+                        help='job launcher')
     parser.add_argument('--local_rank', type=int, default=0)
 
     args = parser.parse_args()
@@ -181,37 +181,33 @@ def main():
     target_transform = transforms.Lambda(
         lambda x: torch.tensor(np.array(x), dtype=torch.long))
 
-    train_set = CamVid(
-        'data/CamVid',
-        img_folder='train',
-        mask_folder='train_labels',
-        transform=transform,
-        target_transform=target_transform)
-
-    valid_set = CamVid(
-        'data/CamVid',
-        img_folder='val',
-        mask_folder='val_labels',
-        transform=transform,
-        target_transform=target_transform)
-
-    train_dataloader = dict(
-        batch_size=3,
-        dataset=train_set,
-        sampler=dict(type='DefaultSampler', shuffle=True),
-        collate_fn=dict(type='default_collate'))
-    val_dataloader = dict(
-        batch_size=3,
-        dataset=valid_set,
-        sampler=dict(type='DefaultSampler', shuffle=False),
-        collate_fn=dict(type='default_collate'))
+    train_set = CamVid('data/CamVid',
+                       img_folder='train',
+                       mask_folder='train_labels',
+                       transform=transform,
+                       target_transform=target_transform)
+
+    valid_set = CamVid('data/CamVid',
+                       img_folder='val',
+                       mask_folder='val_labels',
+                       transform=transform,
+                       target_transform=target_transform)
+
+    train_dataloader = dict(batch_size=3,
+                            dataset=train_set,
+                            sampler=dict(type='DefaultSampler', shuffle=True),
+                            collate_fn=dict(type='default_collate'))
+    val_dataloader = dict(batch_size=3,
+                          dataset=valid_set,
+                          sampler=dict(type='DefaultSampler', shuffle=False),
+                          collate_fn=dict(type='default_collate'))
 
     runner = Runner(
         model=MMDeeplabV3(num_classes),
         work_dir='./work_dir',
         train_dataloader=train_dataloader,
-        optim_wrapper=dict(
-            type=AmpOptimWrapper, optimizer=dict(type=AdamW, lr=2e-4)),
+        optim_wrapper=dict(type=AmpOptimWrapper,
+                           optimizer=dict(type=AdamW, lr=2e-4)),
         train_cfg=dict(by_epoch=True, max_epochs=10, val_interval=10),
         val_dataloader=val_dataloader,
         val_cfg=dict(),

examples/test_time_augmentation.py

@@ -28,15 +28,14 @@ def _merge_single_sample(self, data_samples):
     cfg.work_dir = 'work_dirs/resnet50_8xb16_cifar10'
     cfg.model = dict(type='ClsTTAModel', module=cfg.model)
     test_pipeline = deepcopy(cfg.test_dataloader.dataset.pipeline)
-    flip_tta = dict(
-        type='TestTimeAug',
-        transforms=[
-            [
-                dict(type='RandomFlip', prob=1.),
-                dict(type='RandomFlip', prob=0.)
-            ],
-            [test_pipeline[-1]],
-        ])
+    flip_tta = dict(type='TestTimeAug',
+                    transforms=[
+                        [
+                            dict(type='RandomFlip', prob=1.),
+                            dict(type='RandomFlip', prob=0.)
+                        ],
+                        [test_pipeline[-1]],
+                    ])
     # Replace the last transform with `TestTimeAug`
     cfg.test_dataloader.dataset.pipeline[-1] = flip_tta
     cfg.load_from = 'https://download.openmmlab.com/mmclassification/v0' \