Accueil Explorer Aide
Connexion
mengxi
/
mmdetection
2
0
Fork 0
Fichiers Tickets 0 Pull Requests 0 Wiki
Aborescence: 3a73e65767
Branches Tags
mmdetection
/
tests
/
test_engine
/
test_hooks
/
test_pipeline_switch_hook.py

test_pipeline_switch_hook.py 3.4 KB
Historique Raw

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384 
  1. # Copyright (c) OpenMMLab. All rights reserved.
    2. 

  2. from unittest import TestCase
    3. 

  3. from unittest.mock import Mock
    4. 

  4. 

  5. from mmdet.engine.hooks import PipelineSwitchHook
    6. 

  6. 

  7. 

  8. class TestPipelineSwitchHook(TestCase):
    9. 

  9. 

  10.     def test_persistent_workers_on(self):
    11. 

  11.         runner = Mock()
    12. 

  12.         runner.model = Mock()
    13. 

  13.         runner.model.module = Mock()
    14. 

  14.         runner.train_dataloader = Mock()
    15. 

  15.         runner.train_dataloader.persistent_workers = True
    16. 

  16.         runner.train_dataloader._DataLoader__initialized = True
    17. 

  17. 

  18.         stage2 = [dict(type='RandomResize', scale=(1280, 1280))]
    19. 

  19. 

  20.         runner.epoch = 284  # epoch < switch_epoch
    21. 

  21.         hook = PipelineSwitchHook(switch_epoch=285, switch_pipeline=stage2)
    22. 

  22.         hook.before_train_epoch(runner)
    23. 

  23.         self.assertFalse(hook._restart_dataloader)
    24. 

  24.         self.assertTrue(runner.train_dataloader._DataLoader__initialized)
    25. 

  25. 

  26.         runner.epoch = 285  # epoch == switch_epoch
    27. 

  27.         hook.before_train_epoch(runner)
    28. 

  28.         self.assertTrue(hook._restart_dataloader)
    29. 

  29.         self.assertFalse(runner.train_dataloader._DataLoader__initialized)
    30. 

  30.         self.assertTrue(
    31. 

  31.             len(runner.train_dataloader.dataset.pipeline.transforms) == 1)
    32. 

  32. 

  33.         runner.epoch = 286  # epoch > switch_epoch
    34. 

  34.         hook.before_train_epoch(runner)
    35. 

  35.         self.assertTrue(runner.train_dataloader._DataLoader__initialized)
    36. 

  36.         self.assertTrue(
    37. 

  37.             len(runner.train_dataloader.dataset.pipeline.transforms) == 1)
    38. 

  38. 

  39.     def test_persistent_workers_off(self):
    40. 

  40.         runner = Mock()
    41. 

  41.         runner.model = Mock()
    42. 

  42.         runner.train_dataloader = Mock()
    43. 

  43.         runner.train_dataloader.persistent_workers = False
    44. 

  44.         runner.train_dataloader._DataLoader__initialized = True
    45. 

  45. 

  46.         stage2 = [dict(type='RandomResize', scale=(1280, 1280))]
    47. 

  47. 

  48.         runner.epoch = 284  # epoch < switch_epoch
    49. 

  49.         hook = PipelineSwitchHook(switch_epoch=285, switch_pipeline=stage2)
    50. 

  50.         hook.before_train_epoch(runner)
    51. 

  51.         self.assertFalse(hook._restart_dataloader)
    52. 

  52.         self.assertTrue(runner.train_dataloader._DataLoader__initialized)
    53. 

  53. 

  54.         runner.epoch = 285  # epoch == switch_epoch
    55. 

  55.         hook.before_train_epoch(runner)
    56. 

  56.         self.assertFalse(hook._restart_dataloader)
    57. 

  57.         self.assertTrue(runner.train_dataloader._DataLoader__initialized)
    58. 

  58.         self.assertTrue(
    59. 

  59.             len(runner.train_dataloader.dataset.pipeline.transforms) == 1)
    60. 

  60. 

  61.         runner.epoch = 286  # epoch > switch_epoch
    62. 

  62.         hook.before_train_epoch(runner)
    63. 

  63.         self.assertTrue(runner.train_dataloader._DataLoader__initialized)
    64. 

  64.         self.assertTrue(
    65. 

  65.             len(runner.train_dataloader.dataset.pipeline.transforms) == 1)
    66. 

  66. 

  67.     def test_initialize_after_switching(self):
    68. 

  68.         # This simulates the resumption after the switching.
    69. 

  69.         runner = Mock()
    70. 

  70.         runner.model = Mock()
    71. 

  71.         runner.model.module = Mock()
    72. 

  72.         runner.train_dataloader = Mock()
    73. 

  73.         runner.train_dataloader.persistent_workers = True
    74. 

  74.         runner.train_dataloader._DataLoader__initialized = True
    75. 

  75. 

  76.         stage2 = [dict(type='RandomResize', scale=(1280, 1280))]
    77. 

  77. 

  78.         runner.epoch = 286  # epoch > switch_epoch
    79. 

  79.         hook = PipelineSwitchHook(switch_epoch=285, switch_pipeline=stage2)
    80. 

  80.         hook.before_train_epoch(runner)
    81. 

  81.         self.assertTrue(hook._restart_dataloader)
    82. 

  82.         self.assertFalse(runner.train_dataloader._DataLoader__initialized)
    83. 

  83.         self.assertTrue(
    84. 

  84.             len(runner.train_dataloader.dataset.pipeline.transforms) == 1)
    85.