[Feature] Add RecorderHook

examples/attribute_toy_test.py

@@ -0,0 +1,54 @@
+import torch
+from torch import nn
+from torch.utils.data import DataLoader
+
+from mmengine.model import BaseModel
+from mmengine.runner import Runner
+
+
+class ToyModel(BaseModel):
+
+    def __init__(self, data_preprocessor=None):
+        super().__init__(data_preprocessor=data_preprocessor)
+        self.linear1 = nn.Linear(2, 2)
+        self.linear2 = nn.Linear(2, 1)
+
+    def forward(self, inputs, data_samples, mode='tensor'):
+        if isinstance(inputs, list):
+            inputs = torch.stack(inputs)
+        if isinstance(data_samples, list):
+            data_sample = torch.stack(data_samples)
+        outputs = self.linear1(inputs)
+        outputs = self.linear2(outputs)
+
+        if mode == 'tensor':
+            return outputs
+        elif mode == 'loss':
+            loss = (data_sample - outputs).sum()
+            outputs = dict(loss=loss)
+            return outputs
+        elif mode == 'predict':
+            return outputs
+
+
+x = [(torch.ones(2, 2), [torch.ones(2, 1)])]
+# train_dataset = [x, x, x]
+train_dataset = x * 50
+train_dataloader = DataLoader(train_dataset, batch_size=1)
+
+runner = Runner(
+    model=ToyModel(),
+    custom_hooks=[
+        dict(
+            type='RecorderHook',
+            recorders=[
+                dict(type='AttributeRecorder', target='self.linear1.weight')
+            ],
+            save_dir='./work_dir',
+            print_modification=True)
+    ],
+    work_dir='tmp_dir',
+    train_dataloader=train_dataloader,
+    train_cfg=dict(by_epoch=True, max_epochs=10),
+    optim_wrapper=dict(optimizer=dict(type='SGD', lr=0.01)))
+runner.train()

examples/function_toy_test.py

@@ -0,0 +1,54 @@
+import torch
+from torch import nn
+from torch.utils.data import DataLoader
+
+from mmengine.model import BaseModel
+from mmengine.runner import Runner
+
+
+class ToyModel(BaseModel):
+
+    def __init__(self, data_preprocessor=None):
+        super().__init__(data_preprocessor=data_preprocessor)
+        self.linear1 = nn.Linear(2, 2)
+        self.linear2 = nn.Linear(2, 1)
+
+    def forward(self, inputs, data_samples, mode='tensor'):
+        if isinstance(inputs, list):
+            inputs = torch.stack(inputs)
+        if isinstance(data_samples, list):
+            data_sample = torch.stack(data_samples)
+        outputs = self.linear1(inputs)
+        outputs = self.linear2(outputs)
+
+        if mode == 'tensor':
+            return outputs
+        elif mode == 'loss':
+            loss = (data_sample - outputs).sum()
+            outputs = dict(loss=loss)
+            return outputs
+        elif mode == 'predict':
+            return outputs
+
+
+x = [(torch.ones(2, 2), [torch.ones(2, 1)])]
+# train_dataset = [x, x, x]
+train_dataset = x * 50
+train_dataloader = DataLoader(train_dataset, batch_size=2)
+
+runner = Runner(
+    model=ToyModel(),
+    custom_hooks=[
+        dict(
+            type='RecorderHook',
+            recorders=[
+                dict(type='FunctionRecorder', target='outputs', index=[1])
+            ],
+            save_dir='./work_dir',
+            print_modification=True)
+    ],
+    work_dir='tmp_dir',
+    train_dataloader=train_dataloader,
+    train_cfg=dict(by_epoch=True, max_epochs=10),
+    optim_wrapper=dict(optimizer=dict(type='SGD', lr=0.01)))
+runner.train()

examples/recorder_hook_test.py

@@ -0,0 +1,100 @@
+import torch.nn.functional as F
+import torchvision
+import torchvision.transforms as transforms
+from torch.optim import SGD
+from torch.utils.data import DataLoader
+
+from mmengine.evaluator import BaseMetric
+from mmengine.model import BaseModel
+from mmengine.runner import Runner
+
+
+class MMResNet50(BaseModel):
+
+    def __init__(self):
+        super().__init__()
+        self.resnet = torchvision.models.resnet50()
+
+    def forward(self, imgs, labels, mode):
+        x = self.resnet(imgs)
+        if mode == 'loss':
+            return {'loss': F.cross_entropy(x, labels)}
+        elif mode == 'predict':
+            return x, labels
+
+
+class Accuracy(BaseMetric):
+
+    def process(self, data_batch, data_samples):
+        score, gt = data_samples
+        self.results.append({
+            'batch_size': len(gt),
+            'correct': (score.argmax(dim=1) == gt).sum().cpu(),
+        })
+
+    def compute_metrics(self, results):
+        total_correct = sum(item['correct'] for item in results)
+        total_size = sum(item['batch_size'] for item in results)
+        return dict(accuracy=100 * total_correct / total_size)
+
+
+norm_cfg = dict(mean=[0.491, 0.482, 0.447], std=[0.202, 0.199, 0.201])
+train_dataloader = DataLoader(
+    batch_size=32,
+    shuffle=True,
+    dataset=torchvision.datasets.CIFAR10(
+        'data/cifar10',
+        train=True,
+        download=True,
+        transform=transforms.Compose([
+            transforms.RandomCrop(32, padding=4),
+            transforms.RandomHorizontalFlip(),
+            transforms.ToTensor(),
+            transforms.Normalize(**norm_cfg)
+        ])))
+
+val_dataloader = DataLoader(
+    batch_size=32,
+    shuffle=False,
+    dataset=torchvision.datasets.CIFAR10(
+        'data/cifar10',
+        train=False,
+        download=True,
+        transform=transforms.Compose(
+            [transforms.ToTensor(),
+             transforms.Normalize(**norm_cfg)])))
+
+# default_hooks = dict(logger=dict(type='LoggerHook', interval=20))
+
+runner = Runner(
+    # custom_hooks=[
+    #     dict(
+    #         type='RecorderHook',
+    #         recorders=[dict(type='FunctionRecorder', target='x')],
+    #         save_dir='./work_dir',
+    #         print_modification=True)
+    # ],
+    custom_hooks=[
+        dict(
+            type='RecorderHook',
+            recorders=[
+                dict(
+                    model='resnet',
+                    method='_forward_impl',
+                    type='FunctionRecorder',
+                    target='x',
+                    index=[0, 1, 2])
+            ],
+            save_dir='./work_dir',
+            print_modification=True)
+    ],
+    model=MMResNet50(),
+    work_dir='./work_dir',
+    train_dataloader=train_dataloader,
+    optim_wrapper=dict(optimizer=dict(type=SGD, lr=0.001, momentum=0.9)),
+    train_cfg=dict(by_epoch=True, max_epochs=1, val_interval=1),
+    val_dataloader=val_dataloader,
+    val_cfg=dict(),
+    val_evaluator=dict(type=Accuracy),
+)
+runner.train()

mmengine/hooks/__init__.py

@@ -9,6 +9,7 @@
 from .naive_visualization_hook import NaiveVisualizationHook
 from .param_scheduler_hook import ParamSchedulerHook
 from .profiler_hook import NPUProfilerHook, ProfilerHook
+from .recorder_hook import RecorderHook
 from .runtime_info_hook import RuntimeInfoHook
 from .sampler_seed_hook import DistSamplerSeedHook
 from .sync_buffer_hook import SyncBuffersHook
@@ -18,5 +19,5 @@
     'Hook', 'IterTimerHook', 'DistSamplerSeedHook', 'ParamSchedulerHook',
     'SyncBuffersHook', 'EmptyCacheHook', 'CheckpointHook', 'LoggerHook',
     'NaiveVisualizationHook', 'EMAHook', 'RuntimeInfoHook', 'ProfilerHook',
-    'PrepareTTAHook', 'NPUProfilerHook', 'EarlyStoppingHook'
+    'PrepareTTAHook', 'NPUProfilerHook', 'EarlyStoppingHook', 'RecorderHook'
 ]