|
|
| """
|
| 基础测试
|
| 测试项目的基本功能
|
| """
|
|
|
| import os
|
| import sys
|
| import torch
|
| import torch.nn as nn
|
| import numpy as np
|
|
|
|
|
| sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(__file__))))
|
|
|
| from src.models.unet_light import UNetLight, TimestepEmbedding, ResNetBlock
|
| from src.models.attention import MemoryEfficientAttention
|
| from src.models.diffusion import DiffusionProcess
|
|
|
|
|
| def test_timestep_embedding():
|
| """测试时间步嵌入"""
|
| print("测试时间步嵌入...")
|
|
|
| embedding_dim = 256
|
| time_embed_dim = 512
|
|
|
| embedder = TimestepEmbedding(embedding_dim, time_embed_dim)
|
|
|
|
|
| timesteps = torch.tensor([100, 200, 300])
|
| embeddings = embedder(timesteps)
|
|
|
| assert embeddings.shape == (3, time_embed_dim)
|
| print(f" 形状正确: {embeddings.shape}")
|
|
|
| return embedder
|
|
|
|
|
| def test_resnet_block():
|
| """测试残差块"""
|
| print("测试残差块...")
|
|
|
| in_channels = 64
|
| out_channels = 128
|
| time_embed_dim = 256
|
|
|
| block = ResNetBlock(in_channels, out_channels, time_embed_dim)
|
|
|
|
|
| x = torch.randn(2, in_channels, 32, 32)
|
| time_emb = torch.randn(2, time_embed_dim)
|
|
|
| output = block(x, time_emb)
|
|
|
| assert output.shape == (2, out_channels, 32, 32)
|
| print(f" 形状正确: {output.shape}")
|
|
|
|
|
| block_same = ResNetBlock(in_channels, in_channels, time_embed_dim)
|
| output_same = block_same(x, time_emb)
|
|
|
| assert output_same.shape == x.shape
|
| print(f" 跳跃连接正确")
|
|
|
| return block
|
|
|
|
|
| def test_attention():
|
| """测试注意力机制"""
|
| print("测试注意力机制...")
|
|
|
| dim = 256
|
| num_heads = 8
|
|
|
| attention = MemoryEfficientAttention(dim, num_heads)
|
|
|
|
|
| x = torch.randn(2, 16, dim)
|
| output = attention(x)
|
|
|
| assert output.shape == x.shape
|
| print(f" 形状正确: {output.shape}")
|
|
|
| return attention
|
|
|
|
|
| def test_unet_light():
|
| """测试轻量UNet"""
|
| print("测试轻量UNet...")
|
|
|
| config = {
|
| 'model': {
|
| 'in_channels': 4,
|
| 'out_channels': 4,
|
| 'base_channels': 32,
|
| 'channel_mults': [1, 2, 4],
|
| 'num_res_blocks': 1,
|
| 'attention_resolutions': [8],
|
| 'dropout': 0.0,
|
| 'use_checkpoint': False,
|
| 'num_heads': 4,
|
| 'context_dim': 256,
|
| 'use_linear_projection': True,
|
| 'time_embed_dim': 128
|
| }
|
| }
|
|
|
| model = UNetLight(config)
|
|
|
|
|
| batch_size = 2
|
| x = torch.randn(batch_size, 4, 64, 64)
|
| timesteps = torch.randint(0, 1000, (batch_size,))
|
| context = torch.randn(batch_size, 77, 256)
|
|
|
| output = model(x, timesteps, context)
|
|
|
| assert output.shape == x.shape
|
| print(f" 形状正确: {output.shape}")
|
|
|
|
|
| model.enable_gradient_checkpointing()
|
| print(f" 梯度检查点已启用")
|
|
|
| return model
|
|
|
|
|
| def test_diffusion_process():
|
| """测试扩散过程"""
|
| print("测试扩散过程...")
|
|
|
| config = {
|
| 'diffusion': {
|
| 'beta_schedule': 'linear',
|
| 'beta_start': 0.0001,
|
| 'beta_end': 0.02,
|
| 'num_train_timesteps': 100,
|
| 'num_inference_timesteps': 20
|
| }
|
| }
|
|
|
| diffusion = DiffusionProcess(config)
|
|
|
|
|
| x_start = torch.randn(2, 3, 32, 32)
|
| t = torch.randint(0, 100, (2,))
|
|
|
| x_noisy = diffusion.q_sample(x_start, t)
|
|
|
| assert x_noisy.shape == x_start.shape
|
| print(f" 前向扩散形状正确: {x_noisy.shape}")
|
|
|
|
|
| extracted = diffusion.extract(diffusion.sqrt_alphas_cumprod, t, x_start.shape)
|
| assert extracted.shape == (2, 1, 1, 1)
|
| print(f" 参数提取形状正确: {extracted.shape}")
|
|
|
| return diffusion
|
|
|
|
|
| def test_memory_efficiency():
|
| """测试内存效率"""
|
| print("测试内存效率...")
|
|
|
|
|
| config = {
|
| 'model': {
|
| 'in_channels': 4,
|
| 'out_channels': 4,
|
| 'base_channels': 32,
|
| 'channel_mults': [1, 2],
|
| 'num_res_blocks': 1,
|
| 'attention_resolutions': [],
|
| 'dropout': 0.0,
|
| 'use_checkpoint': False,
|
| 'num_heads': 4,
|
| 'context_dim': 256,
|
| 'use_linear_projection': True,
|
| 'time_embed_dim': 128
|
| }
|
| }
|
|
|
| model = UNetLight(config)
|
| model.eval()
|
|
|
| if torch.cuda.is_available():
|
| device = torch.device('cuda')
|
| model = model.to(device)
|
|
|
| print(" GPU内存测试:")
|
|
|
| for batch_size in [1, 2, 4]:
|
|
|
| torch.cuda.empty_cache()
|
|
|
|
|
| initial_memory = torch.cuda.memory_allocated()
|
|
|
|
|
| x = torch.randn(batch_size, 4, 64, 64, device=device)
|
| t = torch.randint(0, 1000, (batch_size,), device=device)
|
| context = torch.randn(batch_size, 77, 256, device=device)
|
|
|
| with torch.no_grad():
|
| _ = model(x, t, context)
|
|
|
|
|
| peak_memory = torch.cuda.max_memory_allocated()
|
| memory_used = (peak_memory - initial_memory) / 1024**3
|
|
|
| print(f" 批次大小 {batch_size}: {memory_used:.2f} GB")
|
|
|
| else:
|
| print(" GPU不可用，跳过内存测试")
|
|
|
| return model
|
|
|
|
|
| def run_all_tests():
|
| """运行所有测试"""
|
| print("=" * 60)
|
| print("运行Lumina基础测试")
|
| print("=" * 60)
|
|
|
| try:
|
|
|
| test_timestep_embedding()
|
| test_resnet_block()
|
| test_attention()
|
| test_unet_light()
|
| test_diffusion_process()
|
| test_memory_efficiency()
|
|
|
| print("\n" + "=" * 60)
|
| print("所有测试通过!")
|
| print("=" * 60)
|
|
|
| return True
|
|
|
| except Exception as e:
|
| print(f"\n测试失败: {e}")
|
| import traceback
|
| traceback.print_exc()
|
| return False
|
|
|
|
|
| if __name__ == "__main__":
|
| success = run_all_tests()
|
| sys.exit(0 if success else 1) |
