[Fix]: fix error and add unit test (#429)

pull/441/head
Mashiro 2022-08-17 19:16:00 +08:00 committed by GitHub
parent f5cb45dc33
commit a706bbc018
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 15 additions and 1 deletions

View File

@ -34,7 +34,7 @@ class RuntimeInfoHook(Hook):
runner.message_hub.update_info('iter', runner.iter)
runner.message_hub.update_info('max_epochs', runner.max_epochs)
runner.message_hub.update_info('max_iters', runner.max_iters)
if hasattr(runner.train_dataloader.dataset, 'dataset_meta'):
if hasattr(runner.train_dataloader.dataset, 'metainfo'):
runner.message_hub.update_info(
'dataset_meta', runner.train_dataloader.dataset.metainfo)

View File

@ -15,18 +15,32 @@ class TestRuntimeInfoHook(TestCase):
def test_before_train(self):
message_hub = MessageHub.get_instance(
'runtime_info_hook_test_before_train')
class ToyDataset:
...
runner = Mock()
runner.epoch = 7
runner.iter = 71
runner.max_epochs = 4
runner.max_iters = 40
runner.message_hub = message_hub
runner.train_dataloader.dataset = ToyDataset()
hook = RuntimeInfoHook()
hook.before_train(runner)
self.assertEqual(message_hub.get_info('epoch'), 7)
self.assertEqual(message_hub.get_info('iter'), 71)
self.assertEqual(message_hub.get_info('max_epochs'), 4)
self.assertEqual(message_hub.get_info('max_iters'), 40)
with self.assertRaisesRegex(KeyError, 'dataset_meta is not found'):
message_hub.get_info('dataset_meta')
class ToyDatasetWithMeta:
metainfo = dict()
runner.train_dataloader.dataset = ToyDatasetWithMeta()
hook.before_train(runner)
self.assertEqual(message_hub.get_info('dataset_meta'), dict())
def test_before_train_epoch(self):
message_hub = MessageHub.get_instance(