| | import copy |
| | import os |
| | import shutil |
| | import tempfile |
| | import unittest |
| |
|
| | import torch |
| | from modelscope import snapshot_download |
| |
|
| | from swift import ResTuningConfig, Swift, SwiftModel |
| |
|
| |
|
| | class TestSwiftResTuning(unittest.TestCase): |
| |
|
| | def setUp(self): |
| | print(('Testing %s.%s' % (type(self).__name__, self._testMethodName))) |
| | self.tmp_dir = tempfile.TemporaryDirectory().name |
| | if not os.path.exists(self.tmp_dir): |
| | os.makedirs(self.tmp_dir) |
| |
|
| | def tearDown(self): |
| | shutil.rmtree(self.tmp_dir) |
| | super().tearDown() |
| |
|
| | def set_random_seed(self, seed=123): |
| | """Set random seed manually to get deterministic results""" |
| | import random |
| | import numpy as np |
| | import torch |
| | random.seed(seed) |
| | np.random.seed(seed) |
| | torch.manual_seed(seed) |
| | torch.cuda.manual_seed(seed) |
| | torch.cuda.manual_seed_all(seed) |
| |
|
| | def model_comparison(self, model, model2): |
| | model_key = list(model.state_dict().keys()) |
| | model2_key = list(model2.state_dict().keys()) |
| | self.assertTrue(model_key == model2_key) |
| | model_val = torch.sum(torch.stack([torch.sum(val) for val in model.state_dict().values()])) |
| | model2_val = torch.sum(torch.stack([torch.sum(val) for val in model2.state_dict().values()])) |
| | self.assertTrue(torch.isclose(model_val, model2_val)) |
| |
|
| | def test_swift_restuning_vit(self): |
| | model_dir = snapshot_download('AI-ModelScope/vit-base-patch16-224') |
| | from transformers import AutoModelForImageClassification |
| | model = AutoModelForImageClassification.from_pretrained(model_dir) |
| | model_swift_1 = copy.deepcopy(model) |
| | model_swift_2 = copy.deepcopy(model) |
| | result_origin = model(torch.ones((1, 3, 224, 224))).logits |
| | print(f'test_swift_restuning_vit result_origin shape: {result_origin.shape}, ' |
| | f'result_origin sum: {torch.sum(result_origin)}') |
| |
|
| | |
| | self.set_random_seed() |
| | restuning_config_1 = ResTuningConfig( |
| | dims=768, |
| | root_modules=r'.*vit.encoder.layer.0$', |
| | stem_modules=r'.*vit.encoder.layer\.\d+$', |
| | target_modules=r'.*vit.layernorm', |
| | target_modules_hook='input', |
| | tuner_cfg='res_adapter', |
| | ) |
| | model_swift_1 = Swift.prepare_model(model_swift_1, config=restuning_config_1) |
| | self.assertTrue(isinstance(model_swift_1, SwiftModel)) |
| | print(model_swift_1.get_trainable_parameters()) |
| | result_swift_1 = model_swift_1(torch.ones((1, 3, 224, 224))).logits |
| | print(f'test_swift_restuning_vit result_swift_1 shape: {result_swift_1.shape}, ' |
| | f'result_swift_1 sum: {torch.sum(result_swift_1)}') |
| |
|
| | |
| | self.set_random_seed() |
| | restuning_config_2 = ResTuningConfig( |
| | dims=768, |
| | root_modules=r'.*vit.encoder.layer.0$', |
| | stem_modules=r'.*vit.encoder.layer\.\d+$', |
| | target_modules=r'.*vit.encoder', |
| | target_modules_hook='output', |
| | target_hidden_pos='last_hidden_state', |
| | tuner_cfg='res_adapter', |
| | ) |
| | model_swift_2 = Swift.prepare_model(model_swift_2, config=restuning_config_2) |
| | self.assertTrue(isinstance(model_swift_2, SwiftModel)) |
| | print(model_swift_2.get_trainable_parameters()) |
| | result_swift_2 = model_swift_2(torch.ones((1, 3, 224, 224))).logits |
| | print(f'test_swift_restuning_vit result_swift_2 shape: {result_swift_2.shape}, ' |
| | f'result_swift_2 sum: {torch.sum(result_swift_2)}') |
| |
|
| | self.assertTrue(all(torch.isclose(result_swift_1, result_swift_2).flatten())) |
| |
|
| | model_swift_1.save_pretrained(self.tmp_dir) |
| | self.assertTrue(os.path.exists(os.path.join(self.tmp_dir, 'default'))) |
| | model_loaded = Swift.from_pretrained(model, self.tmp_dir) |
| | self.model_comparison(model_swift_1, model_loaded) |
| |
|
| | def test_swift_restuning_diffusers_sd(self): |
| | model_dir = snapshot_download('AI-ModelScope/stable-diffusion-v1-5') |
| | from diffusers import UNet2DConditionModel |
| | model = UNet2DConditionModel.from_pretrained(model_dir, subfolder='unet') |
| | model.requires_grad_(False) |
| | model2 = copy.deepcopy(model) |
| | self.set_random_seed() |
| | input_data = { |
| | 'sample': torch.ones((1, 4, 64, 64)), |
| | 'timestep': 10, |
| | 'encoder_hidden_states': torch.ones((1, 77, 768)) |
| | } |
| | result_origin = model(**input_data).sample |
| | print(f'test_swift_restuning_diffusers_sd result_origin shape: {result_origin.shape}, ' |
| | f'result_origin sum: {torch.sum(result_origin)}') |
| |
|
| | self.set_random_seed() |
| | restuning_config = ResTuningConfig( |
| | dims=[1280, 1280, 1280, 640, 320], |
| | root_modules='mid_block', |
| | stem_modules=['mid_block', 'up_blocks.0', 'up_blocks.1', 'up_blocks.2', 'up_blocks.3'], |
| | target_modules='conv_norm_out', |
| | tuner_cfg='res_group_adapter', |
| | use_upsample=True, |
| | upsample_out_channels=[1280, 1280, 640, 320, None], |
| | zero_init_last=True) |
| |
|
| | model = Swift.prepare_model(model, config=restuning_config) |
| | self.assertTrue(isinstance(model, SwiftModel)) |
| | print(model.get_trainable_parameters()) |
| |
|
| | result = model(**input_data).sample |
| | print(f'test_swift_restuning_diffusers_sd result shape: {result.shape}, result sum: {torch.sum(result)}') |
| | model.save_pretrained(self.tmp_dir) |
| | self.assertTrue(os.path.exists(os.path.join(self.tmp_dir, 'default'))) |
| | model2 = Swift.from_pretrained(model2, self.tmp_dir) |
| | self.model_comparison(model, model2) |
| |
|
| |
|
| | if __name__ == '__main__': |
| | unittest.main() |
| |
|
