modelscope
/
ModelScope
Not watched
1
0
Fork 0
Code
Releases 0 Wiki evaluate Activity
Issues 0 Pull Requests 0 Datasets Model Cloudbrain HPC
You can not select more than 25 topics Topics must start with a chinese character,a letter or number, can include dashes ('-') and can be up to 35 characters long.
Tree: 930d55d9ad
Branches Tags
${ item.name }
Create branch ${ searchTerm }
from '930d55d9ad'
${ noResults }
ModelScope/tests/trainers/test_trainer.py

450 lines
15 kB
Raw Blame History 

  1. # Copyright (c) Alibaba, Inc. and its affiliates.

  2. import os

  3. import shutil

  4. import tempfile

  5. import unittest

  6. 

  7. import json

  8. import numpy as np

  9. import torch

  10. from torch import nn

  11. from torch.optim import SGD

  12. from torch.optim.lr_scheduler import StepLR

  13. from torch.utils.data import IterableDataset

  14. 

  15. from modelscope.metainfo import Metrics, Trainers

  16. from modelscope.metrics.builder import MetricKeys

  17. from modelscope.models.base import Model

  18. from modelscope.trainers import build_trainer

  19. from modelscope.utils.constant import LogKeys, ModeKeys, ModelFile

  20. from modelscope.utils.test_utils import create_dummy_test_dataset, test_level

  21. 

  22. 

  23. class DummyIterableDataset(IterableDataset):

  24. 

  25.     def __iter__(self):

  26.         feat = np.random.random(size=(5, )).astype(np.float32)

  27.         labels = np.random.randint(0, 4, (1, ))

  28.         iterations = [{'feat': feat, 'labels': labels}] * 500

  29.         return iter(iterations)

  30. 

  31. 

  32. dummy_dataset_small = create_dummy_test_dataset(

  33.     np.random.random(size=(5, )), np.random.randint(0, 4, (1, )), 20)

  34. 

  35. dummy_dataset_big = create_dummy_test_dataset(

  36.     np.random.random(size=(5, )), np.random.randint(0, 4, (1, )), 40)

  37. 

  38. 

  39. class DummyModel(nn.Module, Model):

  40. 

  41.     def __init__(self):

  42.         super().__init__()

  43.         self.linear = nn.Linear(5, 4)

  44.         self.bn = nn.BatchNorm1d(4)

  45. 

  46.     def forward(self, feat, labels):

  47.         x = self.linear(feat)

  48. 

  49.         x = self.bn(x)

  50.         loss = torch.sum(x)

  51.         return dict(logits=x, loss=loss)

  52. 

  53. 

  54. class TrainerTest(unittest.TestCase):

  55. 

  56.     def setUp(self):

  57.         print(('Testing %s.%s' % (type(self).__name__, self._testMethodName)))

  58.         self.tmp_dir = tempfile.TemporaryDirectory().name

  59.         if not os.path.exists(self.tmp_dir):

  60.             os.makedirs(self.tmp_dir)

  61. 

  62.     def tearDown(self):

  63.         super().tearDown()

  64.         shutil.rmtree(self.tmp_dir)

  65. 

  66.     @unittest.skipUnless(test_level() >= 1, 'skip test in current test level')

  67.     def test_train_0(self):

  68.         json_cfg = {

  69.             'train': {

  70.                 'work_dir':

  71.                 self.tmp_dir,

  72.                 'dataloader': {

  73.                     'batch_size_per_gpu': 2,

  74.                     'workers_per_gpu': 1

  75.                 },

  76.                 'optimizer': {

  77.                     'type': 'SGD',

  78.                     'lr': 0.01,

  79.                     'options': {

  80.                         'grad_clip': {

  81.                             'max_norm': 2.0

  82.                         }

  83.                     }

  84.                 },

  85.                 'lr_scheduler': {

  86.                     'type': 'StepLR',

  87.                     'step_size': 2,

  88.                     'options': {

  89.                         'warmup': {

  90.                             'type': 'LinearWarmup',

  91.                             'warmup_iters': 2

  92.                         }

  93.                     }

  94.                 },

  95.                 'hooks': [{

  96.                     'type': 'CheckpointHook',

  97.                     'interval': 1

  98.                 }, {

  99.                     'type': 'TextLoggerHook',

  100.                     'interval': 1

  101.                 }, {

  102.                     'type': 'IterTimerHook'

  103.                 }, {

  104.                     'type': 'EvaluationHook',

  105.                     'interval': 1

  106.                 }]

  107.             },

  108.             'evaluation': {

  109.                 'dataloader': {

  110.                     'batch_size_per_gpu': 2,

  111.                     'workers_per_gpu': 1,

  112.                     'shuffle': False

  113.                 },

  114.                 'metrics': [Metrics.seq_cls_metric]

  115.             }

  116.         }

  117.         config_path = os.path.join(self.tmp_dir, ModelFile.CONFIGURATION)

  118.         with open(config_path, 'w') as f:

  119.             json.dump(json_cfg, f)

  120. 

  121.         trainer_name = Trainers.default

  122.         kwargs = dict(

  123.             cfg_file=config_path,

  124.             model=DummyModel(),

  125.             data_collator=None,

  126.             train_dataset=dummy_dataset_small,

  127.             eval_dataset=dummy_dataset_small,

  128.             max_epochs=3,

  129.             device='cpu')

  130. 

  131.         trainer = build_trainer(trainer_name, kwargs)

  132.         trainer.train()

  133.         results_files = os.listdir(self.tmp_dir)

  134. 

  135.         self.assertIn(f'{trainer.timestamp}.log.json', results_files)

  136.         self.assertIn(f'{LogKeys.EPOCH}_1.pth', results_files)

  137.         self.assertIn(f'{LogKeys.EPOCH}_2.pth', results_files)

  138.         self.assertIn(f'{LogKeys.EPOCH}_3.pth', results_files)

  139. 

  140.     @unittest.skipUnless(test_level() >= 1, 'skip test in current test level')

  141.     def test_train_1(self):

  142.         json_cfg = {

  143.             'train': {

  144.                 'work_dir':

  145.                 self.tmp_dir,

  146.                 'dataloader': {

  147.                     'batch_size_per_gpu': 2,

  148.                     'workers_per_gpu': 1

  149.                 },

  150.                 'hooks': [{

  151.                     'type': 'CheckpointHook',

  152.                     'interval': 1

  153.                 }, {

  154.                     'type': 'TextLoggerHook',

  155.                     'interval': 1

  156.                 }, {

  157.                     'type': 'IterTimerHook'

  158.                 }, {

  159.                     'type': 'EvaluationHook',

  160.                     'interval': 1

  161.                 }]

  162.             },

  163.             'evaluation': {

  164.                 'dataloader': {

  165.                     'batch_size_per_gpu': 2,

  166.                     'workers_per_gpu': 1,

  167.                     'shuffle': False

  168.                 },

  169.                 'metrics': [Metrics.seq_cls_metric]

  170.             }

  171.         }

  172. 

  173.         config_path = os.path.join(self.tmp_dir, ModelFile.CONFIGURATION)

  174.         with open(config_path, 'w') as f:

  175.             json.dump(json_cfg, f)

  176. 

  177.         model = DummyModel()

  178.         optimmizer = SGD(model.parameters(), lr=0.01)

  179.         lr_scheduler = StepLR(optimmizer, 2)

  180.         trainer_name = Trainers.default

  181.         kwargs = dict(

  182.             cfg_file=config_path,

  183.             model=model,

  184.             data_collator=None,

  185.             train_dataset=dummy_dataset_small,

  186.             eval_dataset=dummy_dataset_small,

  187.             optimizers=(optimmizer, lr_scheduler),

  188.             max_epochs=3,

  189.             device='cpu')

  190. 

  191.         trainer = build_trainer(trainer_name, kwargs)

  192.         trainer.train()

  193.         results_files = os.listdir(self.tmp_dir)

  194. 

  195.         self.assertIn(f'{trainer.timestamp}.log.json', results_files)

  196.         self.assertIn(f'{LogKeys.EPOCH}_1.pth', results_files)

  197.         self.assertIn(f'{LogKeys.EPOCH}_2.pth', results_files)

  198.         self.assertIn(f'{LogKeys.EPOCH}_3.pth', results_files)

  199. 

  200.     @unittest.skipUnless(test_level() >= 1, 'skip test in current test level')

  201.     def test_train_with_default_config(self):

  202.         json_cfg = {

  203.             'train': {

  204.                 'work_dir': self.tmp_dir,

  205.                 'dataloader': {

  206.                     'batch_size_per_gpu': 2,

  207.                     'workers_per_gpu': 1

  208.                 },

  209.                 'hooks': [{

  210.                     'type': 'EvaluationHook',

  211.                     'interval': 1

  212.                 }]

  213.             },

  214.             'evaluation': {

  215.                 'dataloader': {

  216.                     'batch_size_per_gpu': 2,

  217.                     'workers_per_gpu': 1,

  218.                     'shuffle': False

  219.                 },

  220.                 'metrics': [Metrics.seq_cls_metric]

  221.             }

  222.         }

  223. 

  224.         config_path = os.path.join(self.tmp_dir, ModelFile.CONFIGURATION)

  225.         with open(config_path, 'w') as f:

  226.             json.dump(json_cfg, f)

  227. 

  228.         model = DummyModel()

  229.         optimmizer = SGD(model.parameters(), lr=0.01)

  230.         lr_scheduler = StepLR(optimmizer, 2)

  231.         trainer_name = Trainers.default

  232.         kwargs = dict(

  233.             cfg_file=config_path,

  234.             model=model,

  235.             data_collator=None,

  236.             train_dataset=dummy_dataset_big,

  237.             eval_dataset=dummy_dataset_small,

  238.             optimizers=(optimmizer, lr_scheduler),

  239.             max_epochs=3,

  240.             device='cpu')

  241. 

  242.         trainer = build_trainer(trainer_name, kwargs)

  243.         trainer.train()

  244.         results_files = os.listdir(self.tmp_dir)

  245. 

  246.         json_file = os.path.join(self.tmp_dir, f'{trainer.timestamp}.log.json')

  247.         with open(json_file, 'r') as f:

  248.             lines = [i.strip() for i in f.readlines()]

  249.         self.assertDictContainsSubset(

  250.             {

  251.                 LogKeys.MODE: ModeKeys.TRAIN,

  252.                 LogKeys.EPOCH: 1,

  253.                 LogKeys.ITER: 10,

  254.                 LogKeys.LR: 0.01

  255.             }, json.loads(lines[0]))

  256.         self.assertDictContainsSubset(

  257.             {

  258.                 LogKeys.MODE: ModeKeys.TRAIN,

  259.                 LogKeys.EPOCH: 1,

  260.                 LogKeys.ITER: 20,

  261.                 LogKeys.LR: 0.01

  262.             }, json.loads(lines[1]))

  263.         self.assertDictContainsSubset(

  264.             {

  265.                 LogKeys.MODE: ModeKeys.EVAL,

  266.                 LogKeys.EPOCH: 1,

  267.                 LogKeys.ITER: 20

  268.             }, json.loads(lines[2]))

  269.         self.assertDictContainsSubset(

  270.             {

  271.                 LogKeys.MODE: ModeKeys.TRAIN,

  272.                 LogKeys.EPOCH: 2,

  273.                 LogKeys.ITER: 10,

  274.                 LogKeys.LR: 0.01

  275.             }, json.loads(lines[3]))

  276.         self.assertDictContainsSubset(

  277.             {

  278.                 LogKeys.MODE: ModeKeys.TRAIN,

  279.                 LogKeys.EPOCH: 2,

  280.                 LogKeys.ITER: 20,

  281.                 LogKeys.LR: 0.01

  282.             }, json.loads(lines[4]))

  283.         self.assertDictContainsSubset(

  284.             {

  285.                 LogKeys.MODE: ModeKeys.EVAL,

  286.                 LogKeys.EPOCH: 2,

  287.                 LogKeys.ITER: 20

  288.             }, json.loads(lines[5]))

  289.         self.assertDictContainsSubset(

  290.             {

  291.                 LogKeys.MODE: ModeKeys.TRAIN,

  292.                 LogKeys.EPOCH: 3,

  293.                 LogKeys.ITER: 10,

  294.                 LogKeys.LR: 0.001

  295.             }, json.loads(lines[6]))

  296.         self.assertDictContainsSubset(

  297.             {

  298.                 LogKeys.MODE: ModeKeys.TRAIN,

  299.                 LogKeys.EPOCH: 3,

  300.                 LogKeys.ITER: 20,

  301.                 LogKeys.LR: 0.001

  302.             }, json.loads(lines[7]))

  303.         self.assertDictContainsSubset(

  304.             {

  305.                 LogKeys.MODE: ModeKeys.EVAL,

  306.                 LogKeys.EPOCH: 3,

  307.                 LogKeys.ITER: 20

  308.             }, json.loads(lines[8]))

  309.         self.assertIn(f'{LogKeys.EPOCH}_1.pth', results_files)

  310.         self.assertIn(f'{LogKeys.EPOCH}_2.pth', results_files)

  311.         self.assertIn(f'{LogKeys.EPOCH}_3.pth', results_files)

  312.         for i in [0, 1, 3, 4, 6, 7]:

  313.             self.assertIn(LogKeys.DATA_LOAD_TIME, lines[i])

  314.             self.assertIn(LogKeys.ITER_TIME, lines[i])

  315.         for i in [2, 5, 8]:

  316.             self.assertIn(MetricKeys.ACCURACY, lines[i])

  317. 

  318.     @unittest.skipUnless(test_level() >= 1, 'skip test in current test level')

  319.     def test_train_with_iters_per_epoch(self):

  320.         json_cfg = {

  321.             'train': {

  322.                 'work_dir': self.tmp_dir,

  323.                 'dataloader': {

  324.                     'batch_size_per_gpu': 2,

  325.                     'workers_per_gpu': 1

  326.                 },

  327.                 'hooks': [{

  328.                     'type': 'EvaluationHook',

  329.                     'interval': 1

  330.                 }]

  331.             },

  332.             'evaluation': {

  333.                 'dataloader': {

  334.                     'batch_size_per_gpu': 2,

  335.                     'workers_per_gpu': 1,

  336.                     'shuffle': False

  337.                 },

  338.                 'metrics': [Metrics.seq_cls_metric]

  339.             }

  340.         }

  341.         config_path = os.path.join(self.tmp_dir, ModelFile.CONFIGURATION)

  342.         with open(config_path, 'w') as f:

  343.             json.dump(json_cfg, f)

  344. 

  345.         model = DummyModel()

  346.         optimmizer = SGD(model.parameters(), lr=0.01)

  347.         lr_scheduler = StepLR(optimmizer, 2)

  348.         trainer_name = Trainers.default

  349.         kwargs = dict(

  350.             cfg_file=config_path,

  351.             model=model,

  352.             data_collator=None,

  353.             optimizers=(optimmizer, lr_scheduler),

  354.             train_dataset=DummyIterableDataset(),

  355.             eval_dataset=DummyIterableDataset(),

  356.             train_iters_per_epoch=20,

  357.             val_iters_per_epoch=10,

  358.             max_epochs=3,

  359.             device='cpu')

  360. 

  361.         trainer = build_trainer(trainer_name, kwargs)

  362.         trainer.train()

  363.         results_files = os.listdir(self.tmp_dir)

  364.         json_file = os.path.join(self.tmp_dir, f'{trainer.timestamp}.log.json')

  365.         with open(json_file, 'r') as f:

  366.             lines = [i.strip() for i in f.readlines()]

  367.         self.assertDictContainsSubset(

  368.             {

  369.                 LogKeys.MODE: ModeKeys.TRAIN,

  370.                 LogKeys.EPOCH: 1,

  371.                 LogKeys.ITER: 10,

  372.                 LogKeys.LR: 0.01

  373.             }, json.loads(lines[0]))

  374.         self.assertDictContainsSubset(

  375.             {

  376.                 LogKeys.MODE: ModeKeys.TRAIN,

  377.                 LogKeys.EPOCH: 1,

  378.                 LogKeys.ITER: 20,

  379.                 LogKeys.LR: 0.01

  380.             }, json.loads(lines[1]))

  381.         self.assertDictContainsSubset(

  382.             {

  383.                 LogKeys.MODE: ModeKeys.EVAL,

  384.                 LogKeys.EPOCH: 1,

  385.                 LogKeys.ITER: 10

  386.             }, json.loads(lines[2]))

  387.         self.assertDictContainsSubset(

  388.             {

  389.                 LogKeys.MODE: ModeKeys.TRAIN,

  390.                 LogKeys.EPOCH: 2,

  391.                 LogKeys.ITER: 10,

  392.                 LogKeys.LR: 0.01

  393.             }, json.loads(lines[3]))

  394.         self.assertDictContainsSubset(

  395.             {

  396.                 LogKeys.MODE: ModeKeys.TRAIN,

  397.                 LogKeys.EPOCH: 2,

  398.                 LogKeys.ITER: 20,

  399.                 LogKeys.LR: 0.01

  400.             }, json.loads(lines[4]))

  401.         self.assertDictContainsSubset(

  402.             {

  403.                 LogKeys.MODE: ModeKeys.EVAL,

  404.                 LogKeys.EPOCH: 2,

  405.                 LogKeys.ITER: 10

  406.             }, json.loads(lines[5]))

  407.         self.assertDictContainsSubset(

  408.             {

  409.                 LogKeys.MODE: ModeKeys.TRAIN,

  410.                 LogKeys.EPOCH: 3,

  411.                 LogKeys.ITER: 10,

  412.                 LogKeys.LR: 0.001

  413.             }, json.loads(lines[6]))

  414.         self.assertDictContainsSubset(

  415.             {

  416.                 LogKeys.MODE: ModeKeys.TRAIN,

  417.                 LogKeys.EPOCH: 3,

  418.                 LogKeys.ITER: 20,

  419.                 LogKeys.LR: 0.001

  420.             }, json.loads(lines[7]))

  421.         self.assertDictContainsSubset(

  422.             {

  423.                 LogKeys.MODE: ModeKeys.EVAL,

  424.                 LogKeys.EPOCH: 3,

  425.                 LogKeys.ITER: 10

  426.             }, json.loads(lines[8]))

  427.         self.assertIn(f'{LogKeys.EPOCH}_1.pth', results_files)

  428.         self.assertIn(f'{LogKeys.EPOCH}_2.pth', results_files)

  429.         self.assertIn(f'{LogKeys.EPOCH}_3.pth', results_files)

  430.         for i in [0, 1, 3, 4, 6, 7]:

  431.             self.assertIn(LogKeys.DATA_LOAD_TIME, lines[i])

  432.             self.assertIn(LogKeys.ITER_TIME, lines[i])

  433.         for i in [2, 5, 8]:

  434.             self.assertIn(MetricKeys.ACCURACY, lines[i])

  435. 

  436. 

  437. class DummyTrainerTest(unittest.TestCase):

  438. 

  439.     @unittest.skipUnless(test_level() >= 1, 'skip test in current test level')

  440.     def test_dummy(self):

  441.         default_args = dict(cfg_file='configs/examples/train.json')

  442.         trainer = build_trainer('dummy', default_args)

  443. 

  444.         trainer.train()

  445.         trainer.evaluate()

  446. 

  447. 

  448. if __name__ == '__main__':

  449.     unittest.main()