TAI Research

Initial commit: Lumina_Dev_Legacy (archived)

29691f6 5 months ago

14.6 kB

import torch
import torch.nn as nn
from typing import Optional, Dict, Any
import onnx
import onnxruntime as ort
import os


class ModelOptimizer:
    """模型优化器"""
    def __init__(self, model: nn.Module, device: str = 'cuda'):
        self.model = model
        self.device = device
        self.optimized_model = None
    
    def optimize_for_inference(self, use_jit: bool = True, use_cuda_graph: bool = False):
        """优化模型用于推理"""
        self.model.eval()
        
        # 应用一系列优化
        if use_jit:
            self._jit_compile()
        
        if use_cuda_graph and torch.cuda.is_available():
            self._capture_cuda_graph()
        
        # 应用其他优化
        self._apply_inference_optimizations()
        
        return self.optimized_model or self.model
    
    def _jit_compile(self):
        """使用TorchScript编译模型"""
        try:
            # 创建示例输入
            example_input = torch.randn(1, 4, 64, 64, device=self.device)
            example_timestep = torch.tensor([500], device=self.device)
            example_context = torch.randn(1, 77, 768, device=self.device)
            
            # 脚本编译
            scripted_model = torch.jit.trace(
                self.model,
                (example_input, example_timestep, example_context),
                check_trace=False
            )
            
            self.optimized_model = scripted_model
            print("模型已使用TorchScript编译")
        except Exception as e:
            print(f"TorchScript编译失败: {e}")
    
    def _capture_cuda_graph(self):
        """捕获CUDA图（用于重复推理）"""
        if not torch.cuda.is_available():
            return
        
        # 创建静态输入
        static_input = torch.randn(1, 4, 64, 64, device='cuda', dtype=torch.float16)
        static_timestep = torch.tensor([500], device='cuda')
        static_context = torch.randn(1, 77, 768, device='cuda', dtype=torch.float16)
        
        # 预热
        with torch.no_grad():
            for _ in range(3):
                _ = self.model(static_input, static_timestep, static_context)
        
        # 捕获图
        graph = torch.cuda.CUDAGraph()
        with torch.cuda.graph(graph):
            static_output = self.model(static_input, static_timestep, static_context)
        
        # 创建包装函数
        def graph_executor(input_tensor, timestep, context):
            static_input.copy_(input_tensor)
            static_timestep.copy_(timestep)
            static_context.copy_(context)
            graph.replay()
            return static_output.clone()
        
        self.optimized_model = graph_executor
        print("已捕获CUDA图")
    
    def _apply_inference_optimizations(self):
        """应用推理优化"""
        # 设置为评估模式
        self.model.eval()
        
        # 融合操作（如果可用）
        if hasattr(torch, 'compile'):
            try:
                self.model = torch.compile(self.model, mode='max-autotune')
                print("模型已使用torch.compile优化")
            except:
                pass
        
        # 使用半精度
        if self.device == 'cuda':
            self.model.half()
            print("模型已转换为半精度")
    
    def quantize(self, quantization_mode: str = 'dynamic'):
        """量化模型"""
        if quantization_mode == 'dynamic':
            # 动态量化
            quantized_model = torch.quantization.quantize_dynamic(
                self.model,
                {nn.Linear, nn.Conv2d},
                dtype=torch.qint8
            )
            self.optimized_model = quantized_model
            print("模型已动态量化")
        
        elif quantization_mode == 'static':
            # 静态量化需要校准数据
            print("静态量化需要校准数据，暂未实现")
        
        else:
            raise ValueError(f"未知的量化模式: {quantization_mode}")
    
    def prune(self, pruning_rate: float = 0.2):
        """剪枝模型"""
        from torch.nn.utils import prune
        
        # 对线性层和卷积层进行剪枝
        for name, module in self.model.named_modules():
            if isinstance(module, (nn.Linear, nn.Conv2d)):
                prune.l1_unstructured(module, name='weight', amount=pruning_rate)
                prune.remove(module, 'weight')
        
        print(f"模型已剪枝，剪枝率: {pruning_rate}")
    
    def get_model_size(self) -> Dict[str, float]:
        """获取模型大小"""
        # 计算参数量
        total_params = sum(p.numel() for p in self.model.parameters())
        trainable_params = sum(p.numel() for p in self.model.parameters() if p.requires_grad)
        
        # 计算模型大小（MB）
        param_size = 0
        for param in self.model.parameters():
            param_size += param.nelement() * param.element_size()
        
        buffer_size = 0
        for buffer in self.model.buffers():
            buffer_size += buffer.nelement() * buffer.element_size()
        
        size_mb = (param_size + buffer_size) / 1024**2
        
        return {
            'total_params': total_params,
            'trainable_params': trainable_params,
            'size_mb': size_mb
        }


class ONNXExporter:
    """ONNX导出器"""
    def __init__(self, model: nn.Module):
        self.model = model
    
    def export(
        self,
        output_path: str,
        input_shape: tuple = (1, 4, 64, 64),
        opset_version: int = 14,
        dynamic_axes: Optional[Dict] = None
    ):
        """导出为ONNX格式"""
        # 设置为评估模式
        self.model.eval()
        
        # 创建示例输入
        dummy_input = torch.randn(*input_shape)
        dummy_timestep = torch.tensor([500])
        dummy_context = torch.randn(1, 77, 768)
        
        # 默认动态轴
        if dynamic_axes is None:
            dynamic_axes = {
                'input': {0: 'batch_size'},
                'timestep': {0: 'batch_size'},
                'context': {0: 'batch_size'},
                'output': {0: 'batch_size'}
            }
        
        # 导出
        torch.onnx.export(
            self.model,
            (dummy_input, dummy_timestep, dummy_context),
            output_path,
            input_names=['input', 'timestep', 'context'],
            output_names=['output'],
            dynamic_axes=dynamic_axes,
            opset_version=opset_version,
            do_constant_folding=True,
            verbose=False
        )
        
        print(f"模型已导出为ONNX: {output_path}")
        
        # 验证ONNX模型
        self._validate_onnx(output_path)
    
    def _validate_onnx(self, onnx_path: str):
        """验证ONNX模型"""
        try:
            onnx_model = onnx.load(onnx_path)
            onnx.checker.check_model(onnx_model)
            print("ONNX模型验证成功")
        except Exception as e:
            print(f"ONNX模型验证失败: {e}")
    
    def optimize_onnx(self, onnx_path: str, optimized_path: str):
        """优化ONNX模型"""
        try:
            # 使用ONNX Runtime优化
            sess_options = ort.SessionOptions()
            sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
            
            # 创建优化会话
            ort_session = ort.InferenceSession(onnx_path, sess_options)
            
            # 保存优化后的模型
            optimized_model = ort_session.get_model()
            onnx.save(optimized_model, optimized_path)
            
            print(f"ONNX模型已优化: {optimized_path}")
        except Exception as e:
            print(f"ONNX优化失败: {e}")


class MemoryEfficientInference:
    """内存高效推理"""
    def __init__(self, model: nn.Module, chunk_size: int = 32):
        self.model = model
        self.chunk_size = chunk_size
        
    def chunked_inference(self, x: torch.Tensor, t: torch.Tensor, context: torch.Tensor) -> torch.Tensor:
        """分块推理，减少内存使用"""
        B, C, H, W = x.shape
        output = torch.zeros_like(x)
        
        # 分块处理
        for i in range(0, H, self.chunk_size):
            for j in range(0, W, self.chunk_size):
                # 提取块
                chunk = x[:, :, i:i+self.chunk_size, j:j+self.chunk_size]
                
                # 推理
                with torch.no_grad():
                    chunk_output = self.model(chunk, t, context)
                
                # 存储结果
                output[:, :, i:i+self.chunk_size, j:j+self.chunk_size] = chunk_output
        
        return output
    
    def tiled_inference(self, x: torch.Tensor, t: torch.Tensor, context: torch.Tensor, tile_size: int = 512) -> torch.Tensor:
        """平铺推理，用于大图像"""
        B, C, H, W = x.shape
        
        # 如果图像不大，直接推理
        if H <= tile_size and W <= tile_size:
            with torch.no_grad():
                return self.model(x, t, context)
        
        # 计算平铺数量
        n_tiles_h = (H + tile_size - 1) // tile_size
        n_tiles_w = (W + tile_size - 1) // tile_size
        
        output = torch.zeros_like(x)
        
        # 处理每个平铺
        for i in range(n_tiles_h):
            for j in range(n_tiles_w):
                # 计算平铺位置
                h_start = i * tile_size
                w_start = j * tile_size
                h_end = min(h_start + tile_size, H)
                w_end = min(w_start + tile_size, W)
                
                # 提取平铺
                tile = x[:, :, h_start:h_end, w_start:w_end]
                
                # 推理
                with torch.no_grad():
                    tile_output = self.model(tile, t, context)
                
                # 存储结果
                output[:, :, h_start:h_end, w_start:w_end] = tile_output
        
        return output


class InferenceBenchmark:
    """推理基准测试"""
    def __init__(self, model: nn.Module, device: str = 'cuda'):
        self.model = model
        self.device = device
        
    def benchmark(
        self,
        input_shape: tuple = (1, 4, 64, 64),
        num_iterations: int = 100,
        warmup_iterations: int = 10
    ) -> Dict[str, float]:
        """运行基准测试"""
        # 准备输入
        x = torch.randn(*input_shape, device=self.device)
        t = torch.tensor([500], device=self.device)
        context = torch.randn(1, 77, 768, device=self.device)
        
        # 预热
        print("预热...")
        with torch.no_grad():
            for _ in range(warmup_iterations):
                _ = self.model(x, t, context)
        
        # 同步
        if torch.cuda.is_available():
            torch.cuda.synchronize()
        
        # 基准测试
        print("运行基准测试...")
        import time
        
        times = []
        
        for i in range(num_iterations):
            start_time = time.time()
            
            with torch.no_grad():
                _ = self.model(x, t, context)
            
            if torch.cuda.is_available():
                torch.cuda.synchronize()
            
            end_time = time.time()
            times.append(end_time - start_time)
        
        # 统计
        times = torch.tensor(times)
        
        stats = {
            'mean_ms': times.mean().item() * 1000,
            'std_ms': times.std().item() * 1000,
            'min_ms': times.min().item() * 1000,
            'max_ms': times.max().item() * 1000,
            'fps': 1 / times.mean().item(),
            'num_iterations': num_iterations
        }
        
        # 打印结果
        print("\n" + "="*50)
        print("推理基准测试结果:")
        print(f"平均推理时间: {stats['mean_ms']:.2f} ms")
        print(f"标准差: {stats['std_ms']:.2f} ms")
        print(f"最小推理时间: {stats['min_ms']:.2f} ms")
        print(f"最大推理时间: {stats['max_ms']:.2f} ms")
        print(f"FPS: {stats['fps']:.2f}")
        print("="*50)
        
        return stats


def optimize_model_for_p4(model: nn.Module) -> nn.Module:
    """为P4优化模型"""
    optimizer = ModelOptimizer(model)
    
    # 获取模型大小
    size_info = optimizer.get_model_size()
    print(f"优化前模型大小: {size_info['size_mb']:.2f} MB")
    
    # 应用优化
    optimized_model = optimizer.optimize_for_inference(
        use_jit=True,
        use_cuda_graph=False  # P4可能不支持
    )
    
    # 量化（可选）
    if size_info['size_mb'] > 500:  # 如果模型大于500MB，进行量化
        optimizer.quantize('dynamic')
    
    # 获取优化后的模型大小
    size_info_after = optimizer.get_model_size()
    print(f"优化后模型大小: {size_info_after['size_mb']:.2f} MB")
    print(f"压缩比: {size_info['size_mb'] / size_info_after['size_mb']:.2f}x")
    
    return optimized_model


def test_optimization():
    """测试优化"""
    import torch.nn as nn
    
    # 创建模拟模型
    class MockModel(nn.Module):
        def __init__(self):
            super().__init__()
            self.conv1 = nn.Conv2d(4, 64, 3, padding=1)
            self.conv2 = nn.Conv2d(64, 4, 3, padding=1)
        
        def forward(self, x, t, context):
            x = self.conv1(x)
            x = nn.functional.relu(x)
            x = self.conv2(x)
            return x
    
    model = MockModel()
    
    # 测试优化器
    optimizer = ModelOptimizer(model)
    optimized_model = optimizer.optimize_for_inference()
    
    # 测试基准测试
    benchmark = InferenceBenchmark(model)
    stats = benchmark.benchmark(num_iterations=10)
    
    # 测试ONNX导出
    exporter = ONNXExporter(model)
    exporter.export('./test_model.onnx', input_shape=(1, 4, 64, 64))
    
    return optimized_model, stats


if __name__ == '__main__':
    optimized_model, stats = test_optimization()

	import torch
	import torch.nn as nn
	from typing import Optional, Dict, Any
	import onnx
	import onnxruntime as ort
	import os


	class ModelOptimizer:
	"""模型优化器"""
	def __init__(self, model: nn.Module, device: str = 'cuda'):
	self.model = model
	self.device = device
	self.optimized_model = None

	def optimize_for_inference(self, use_jit: bool = True, use_cuda_graph: bool = False):
	"""优化模型用于推理"""
	self.model.eval()

	# 应用一系列优化
	if use_jit:
	self._jit_compile()

	if use_cuda_graph and torch.cuda.is_available():
	self._capture_cuda_graph()

	# 应用其他优化
	self._apply_inference_optimizations()

	return self.optimized_model or self.model

	def _jit_compile(self):
	"""使用TorchScript编译模型"""
	try:
	# 创建示例输入
	example_input = torch.randn(1, 4, 64, 64, device=self.device)
	example_timestep = torch.tensor([500], device=self.device)
	example_context = torch.randn(1, 77, 768, device=self.device)

	# 脚本编译
	scripted_model = torch.jit.trace(
	self.model,
	(example_input, example_timestep, example_context),
	check_trace=False
	)

	self.optimized_model = scripted_model
	print("模型已使用TorchScript编译")
	except Exception as e:
	print(f"TorchScript编译失败: {e}")

	def _capture_cuda_graph(self):
	"""捕获CUDA图（用于重复推理）"""
	if not torch.cuda.is_available():
	return

	# 创建静态输入
	static_input = torch.randn(1, 4, 64, 64, device='cuda', dtype=torch.float16)
	static_timestep = torch.tensor([500], device='cuda')
	static_context = torch.randn(1, 77, 768, device='cuda', dtype=torch.float16)

	# 预热
	with torch.no_grad():
	for _ in range(3):
	_ = self.model(static_input, static_timestep, static_context)

	# 捕获图
	graph = torch.cuda.CUDAGraph()
	with torch.cuda.graph(graph):
	static_output = self.model(static_input, static_timestep, static_context)

	# 创建包装函数
	def graph_executor(input_tensor, timestep, context):
	static_input.copy_(input_tensor)
	static_timestep.copy_(timestep)
	static_context.copy_(context)
	graph.replay()
	return static_output.clone()

	self.optimized_model = graph_executor
	print("已捕获CUDA图")

	def _apply_inference_optimizations(self):
	"""应用推理优化"""
	# 设置为评估模式
	self.model.eval()

	# 融合操作（如果可用）
	if hasattr(torch, 'compile'):
	try:
	self.model = torch.compile(self.model, mode='max-autotune')
	print("模型已使用torch.compile优化")
	except:
	pass

	# 使用半精度
	if self.device == 'cuda':
	self.model.half()
	print("模型已转换为半精度")

	def quantize(self, quantization_mode: str = 'dynamic'):
	"""量化模型"""
	if quantization_mode == 'dynamic':
	# 动态量化
	quantized_model = torch.quantization.quantize_dynamic(
	self.model,
	{nn.Linear, nn.Conv2d},
	dtype=torch.qint8
	)
	self.optimized_model = quantized_model
	print("模型已动态量化")

	elif quantization_mode == 'static':
	# 静态量化需要校准数据
	print("静态量化需要校准数据，暂未实现")

	else:
	raise ValueError(f"未知的量化模式: {quantization_mode}")

	def prune(self, pruning_rate: float = 0.2):
	"""剪枝模型"""
	from torch.nn.utils import prune

	# 对线性层和卷积层进行剪枝
	for name, module in self.model.named_modules():
	if isinstance(module, (nn.Linear, nn.Conv2d)):
	prune.l1_unstructured(module, name='weight', amount=pruning_rate)
	prune.remove(module, 'weight')

	print(f"模型已剪枝，剪枝率: {pruning_rate}")

	def get_model_size(self) -> Dict[str, float]:
	"""获取模型大小"""
	# 计算参数量
	total_params = sum(p.numel() for p in self.model.parameters())
	trainable_params = sum(p.numel() for p in self.model.parameters() if p.requires_grad)

	# 计算模型大小（MB）
	param_size = 0
	for param in self.model.parameters():
	param_size += param.nelement() * param.element_size()

	buffer_size = 0
	for buffer in self.model.buffers():
	buffer_size += buffer.nelement() * buffer.element_size()

	size_mb = (param_size + buffer_size) / 1024**2

	return {
	'total_params': total_params,
	'trainable_params': trainable_params,
	'size_mb': size_mb
	}


	class ONNXExporter:
	"""ONNX导出器"""
	def __init__(self, model: nn.Module):
	self.model = model

	def export(
	self,
	output_path: str,
	input_shape: tuple = (1, 4, 64, 64),
	opset_version: int = 14,
	dynamic_axes: Optional[Dict] = None
	):
	"""导出为ONNX格式"""
	# 设置为评估模式
	self.model.eval()

	# 创建示例输入
	dummy_input = torch.randn(*input_shape)
	dummy_timestep = torch.tensor([500])
	dummy_context = torch.randn(1, 77, 768)

	# 默认动态轴
	if dynamic_axes is None:
	dynamic_axes = {
	'input': {0: 'batch_size'},
	'timestep': {0: 'batch_size'},
	'context': {0: 'batch_size'},
	'output': {0: 'batch_size'}
	}

	# 导出
	torch.onnx.export(
	self.model,
	(dummy_input, dummy_timestep, dummy_context),
	output_path,
	input_names=['input', 'timestep', 'context'],
	output_names=['output'],
	dynamic_axes=dynamic_axes,
	opset_version=opset_version,
	do_constant_folding=True,
	verbose=False
	)

	print(f"模型已导出为ONNX: {output_path}")

	# 验证ONNX模型
	self._validate_onnx(output_path)

	def _validate_onnx(self, onnx_path: str):
	"""验证ONNX模型"""
	try:
	onnx_model = onnx.load(onnx_path)
	onnx.checker.check_model(onnx_model)
	print("ONNX模型验证成功")
	except Exception as e:
	print(f"ONNX模型验证失败: {e}")

	def optimize_onnx(self, onnx_path: str, optimized_path: str):
	"""优化ONNX模型"""
	try:
	# 使用ONNX Runtime优化
	sess_options = ort.SessionOptions()
	sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL

	# 创建优化会话
	ort_session = ort.InferenceSession(onnx_path, sess_options)

	# 保存优化后的模型
	optimized_model = ort_session.get_model()
	onnx.save(optimized_model, optimized_path)

	print(f"ONNX模型已优化: {optimized_path}")
	except Exception as e:
	print(f"ONNX优化失败: {e}")


	class MemoryEfficientInference:
	"""内存高效推理"""
	def __init__(self, model: nn.Module, chunk_size: int = 32):
	self.model = model
	self.chunk_size = chunk_size

	def chunked_inference(self, x: torch.Tensor, t: torch.Tensor, context: torch.Tensor) -> torch.Tensor:
	"""分块推理，减少内存使用"""
	B, C, H, W = x.shape
	output = torch.zeros_like(x)

	# 分块处理
	for i in range(0, H, self.chunk_size):
	for j in range(0, W, self.chunk_size):
	# 提取块
	chunk = x[:, :, i:i+self.chunk_size, j:j+self.chunk_size]

	# 推理
	with torch.no_grad():
	chunk_output = self.model(chunk, t, context)

	# 存储结果
	output[:, :, i:i+self.chunk_size, j:j+self.chunk_size] = chunk_output

	return output

	def tiled_inference(self, x: torch.Tensor, t: torch.Tensor, context: torch.Tensor, tile_size: int = 512) -> torch.Tensor:
	"""平铺推理，用于大图像"""
	B, C, H, W = x.shape

	# 如果图像不大，直接推理
	if H <= tile_size and W <= tile_size:
	with torch.no_grad():
	return self.model(x, t, context)

	# 计算平铺数量
	n_tiles_h = (H + tile_size - 1) // tile_size
	n_tiles_w = (W + tile_size - 1) // tile_size

	output = torch.zeros_like(x)

	# 处理每个平铺
	for i in range(n_tiles_h):
	for j in range(n_tiles_w):
	# 计算平铺位置
	h_start = i * tile_size
	w_start = j * tile_size
	h_end = min(h_start + tile_size, H)
	w_end = min(w_start + tile_size, W)

	# 提取平铺
	tile = x[:, :, h_start:h_end, w_start:w_end]

	# 推理
	with torch.no_grad():
	tile_output = self.model(tile, t, context)

	# 存储结果
	output[:, :, h_start:h_end, w_start:w_end] = tile_output

	return output


	class InferenceBenchmark:
	"""推理基准测试"""
	def __init__(self, model: nn.Module, device: str = 'cuda'):
	self.model = model
	self.device = device

	def benchmark(
	self,
	input_shape: tuple = (1, 4, 64, 64),
	num_iterations: int = 100,
	warmup_iterations: int = 10
	) -> Dict[str, float]:
	"""运行基准测试"""
	# 准备输入
	x = torch.randn(*input_shape, device=self.device)
	t = torch.tensor([500], device=self.device)
	context = torch.randn(1, 77, 768, device=self.device)

	# 预热
	print("预热...")
	with torch.no_grad():
	for _ in range(warmup_iterations):
	_ = self.model(x, t, context)

	# 同步
	if torch.cuda.is_available():
	torch.cuda.synchronize()

	# 基准测试
	print("运行基准测试...")
	import time

	times = []

	for i in range(num_iterations):
	start_time = time.time()

	with torch.no_grad():
	_ = self.model(x, t, context)

	if torch.cuda.is_available():
	torch.cuda.synchronize()

	end_time = time.time()
	times.append(end_time - start_time)

	# 统计
	times = torch.tensor(times)

	stats = {
	'mean_ms': times.mean().item() * 1000,
	'std_ms': times.std().item() * 1000,
	'min_ms': times.min().item() * 1000,
	'max_ms': times.max().item() * 1000,
	'fps': 1 / times.mean().item(),
	'num_iterations': num_iterations
	}

	# 打印结果
	print("\n" + "="*50)
	print("推理基准测试结果:")
	print(f"平均推理时间: {stats['mean_ms']:.2f} ms")
	print(f"标准差: {stats['std_ms']:.2f} ms")
	print(f"最小推理时间: {stats['min_ms']:.2f} ms")
	print(f"最大推理时间: {stats['max_ms']:.2f} ms")
	print(f"FPS: {stats['fps']:.2f}")
	print("="*50)

	return stats


	def optimize_model_for_p4(model: nn.Module) -> nn.Module:
	"""为P4优化模型"""
	optimizer = ModelOptimizer(model)

	# 获取模型大小
	size_info = optimizer.get_model_size()
	print(f"优化前模型大小: {size_info['size_mb']:.2f} MB")

	# 应用优化
	optimized_model = optimizer.optimize_for_inference(
	use_jit=True,
	use_cuda_graph=False # P4可能不支持
	)

	# 量化（可选）
	if size_info['size_mb'] > 500: # 如果模型大于500MB，进行量化
	optimizer.quantize('dynamic')

	# 获取优化后的模型大小
	size_info_after = optimizer.get_model_size()
	print(f"优化后模型大小: {size_info_after['size_mb']:.2f} MB")
	print(f"压缩比: {size_info['size_mb'] / size_info_after['size_mb']:.2f}x")

	return optimized_model


	def test_optimization():
	"""测试优化"""
	import torch.nn as nn

	# 创建模拟模型
	class MockModel(nn.Module):
	def __init__(self):
	super().__init__()
	self.conv1 = nn.Conv2d(4, 64, 3, padding=1)
	self.conv2 = nn.Conv2d(64, 4, 3, padding=1)

	def forward(self, x, t, context):
	x = self.conv1(x)
	x = nn.functional.relu(x)
	x = self.conv2(x)
	return x

	model = MockModel()

	# 测试优化器
	optimizer = ModelOptimizer(model)
	optimized_model = optimizer.optimize_for_inference()

	# 测试基准测试
	benchmark = InferenceBenchmark(model)
	stats = benchmark.benchmark(num_iterations=10)

	# 测试ONNX导出
	exporter = ONNXExporter(model)
	exporter.export('./test_model.onnx', input_shape=(1, 4, 64, 64))

	return optimized_model, stats


	if __name__ == '__main__':
	optimized_model, stats = test_optimization()