wehub-resource-sync
145 lines
6.5 KiB
Python
145 lines
6.5 KiB
Python
import copy
|
|
import os
|
|
import shutil
|
|
import tempfile
|
|
import torch
|
|
import unittest
|
|
from modelscope import snapshot_download
|
|
from transformers.utils import is_torch_npu_available
|
|
|
|
from swift.tuners import SCETuningConfig, Swift
|
|
from swift.tuners.part import PartConfig
|
|
|
|
|
|
def get_npu_or_cpu_device():
|
|
if is_torch_npu_available():
|
|
return torch.device('npu')
|
|
return torch.device('cpu')
|
|
|
|
|
|
def get_diffusers_unet_input(device):
|
|
return {
|
|
'sample': torch.ones((1, 4, 64, 64), device=device),
|
|
'timestep': torch.tensor(10, device=device),
|
|
'encoder_hidden_states': torch.ones((1, 77, 768), device=device)
|
|
}
|
|
|
|
|
|
class TestSCETuning(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 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))
|
|
|
|
@unittest.skip('Legacy test cases')
|
|
def test_scetuning_on_diffusers_v1(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)
|
|
model_check = copy.deepcopy(model)
|
|
device = get_npu_or_cpu_device()
|
|
# module_keys = [key for key, _ in model.named_modules()]
|
|
scetuning_config = SCETuningConfig(
|
|
dims=[320, 320, 320, 320, 640, 640, 640, 1280, 1280, 1280, 1280, 1280],
|
|
tuner_mode='encoder',
|
|
target_modules=[
|
|
'conv_in', 'down_blocks.0.attentions.0', 'down_blocks.0.attentions.1', 'down_blocks.0.downsamplers',
|
|
'down_blocks.1.attentions.0', 'down_blocks.1.attentions.1', 'down_blocks.1.downsamplers',
|
|
'down_blocks.2.attentions.0', 'down_blocks.2.attentions.1', 'down_blocks.2.downsamplers',
|
|
'down_blocks.3.resnets.0', 'down_blocks.3.resnets.1'
|
|
])
|
|
model = Swift.prepare_model(model, config=scetuning_config).to(device)
|
|
model_check = model_check.to(device)
|
|
print(model.get_trainable_parameters())
|
|
input_data = get_diffusers_unet_input(device)
|
|
result = model(**input_data).sample
|
|
print(result.shape)
|
|
model.save_pretrained(self.tmp_dir)
|
|
self.assertTrue(os.path.exists(os.path.join(self.tmp_dir, 'default')))
|
|
model_check = Swift.from_pretrained(model_check, self.tmp_dir).to(device)
|
|
self.model_comparison(model, model_check)
|
|
|
|
@unittest.skip('Legacy test cases')
|
|
def test_scetuning_part_mixin(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)
|
|
model_check = copy.deepcopy(model)
|
|
# module_keys = [key for key, _ in model.named_modules()]
|
|
scetuning_config = SCETuningConfig(
|
|
dims=[320, 320, 320, 320, 640, 640, 640, 1280, 1280, 1280, 1280, 1280],
|
|
tuner_mode='encoder',
|
|
target_modules=[
|
|
'conv_in', 'down_blocks.0.attentions.0', 'down_blocks.0.attentions.1', 'down_blocks.0.downsamplers',
|
|
'down_blocks.1.attentions.0', 'down_blocks.1.attentions.1', 'down_blocks.1.downsamplers',
|
|
'down_blocks.2.attentions.0', 'down_blocks.2.attentions.1', 'down_blocks.2.downsamplers',
|
|
'down_blocks.3.resnets.0', 'down_blocks.3.resnets.1'
|
|
])
|
|
targets = r'.*(to_k|to_v).*'
|
|
part_config = PartConfig(target_modules=targets)
|
|
model = Swift.prepare_model(model, config=scetuning_config)
|
|
model = Swift.prepare_model(model, config={'part': part_config})
|
|
print(model.get_trainable_parameters())
|
|
input_data = {
|
|
'sample': torch.ones((1, 4, 64, 64)),
|
|
'timestep': 10,
|
|
'encoder_hidden_states': torch.ones((1, 77, 768))
|
|
}
|
|
model.set_active_adapters('default')
|
|
model.set_active_adapters('part')
|
|
model.set_active_adapters('default')
|
|
result = model(**input_data).sample
|
|
print(result.shape)
|
|
model.save_pretrained(self.tmp_dir)
|
|
self.assertTrue(os.path.exists(os.path.join(self.tmp_dir, 'default')))
|
|
model_check = Swift.from_pretrained(model_check, self.tmp_dir)
|
|
self.model_comparison(model, model_check)
|
|
|
|
@unittest.skip('Legacy test cases')
|
|
def test_scetuning_on_diffusers_v2(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)
|
|
model_check = copy.deepcopy(model)
|
|
device = get_npu_or_cpu_device()
|
|
# module_keys = [key for key, _ in model.named_modules()]
|
|
scetuning_config = SCETuningConfig(
|
|
dims=[1280, 1280, 1280, 1280, 1280, 640, 640, 640, 320, 320, 320, 320],
|
|
tuner_mode='decoder',
|
|
target_modules=[
|
|
'up_blocks.0.resnets.0', 'up_blocks.0.resnets.1', 'up_blocks.0.resnets.2', 'up_blocks.1.resnets.0',
|
|
'up_blocks.1.resnets.1', 'up_blocks.1.resnets.2', 'up_blocks.2.resnets.0', 'up_blocks.2.resnets.1',
|
|
'up_blocks.2.resnets.2', 'up_blocks.3.resnets.0', 'up_blocks.3.resnets.1', 'up_blocks.3.resnets.2'
|
|
])
|
|
model = Swift.prepare_model(model, config=scetuning_config).to(device)
|
|
model_check = model_check.to(device)
|
|
print(model.get_trainable_parameters())
|
|
input_data = get_diffusers_unet_input(device)
|
|
result = model(**input_data).sample
|
|
print(result.shape)
|
|
model.save_pretrained(self.tmp_dir)
|
|
self.assertTrue(os.path.exists(os.path.join(self.tmp_dir, 'default')))
|
|
model_check = Swift.from_pretrained(model_check, self.tmp_dir).to(device)
|
|
self.model_comparison(model, model_check)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|