|
|
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() |
|
|
|
