增加模型量化支持 (Add Model Quantization Support)

[wowgaoyan177]

问题背景 (Background):
TimeMixer目前没有提供模型量化功能，这限制了其在资源受限设备上的部署能力。量化可以显著减小模型大小并提高推理速度，特别是在边缘设备上。

主要需求 (Requirements):

• 支持INT8和FP16量化
• 提供训练后量化(Post-training quantization)
• 支持量化感知训练(Quantization-aware training)
• 保持性能损失在可接受范围内

技术实现方案 (Implementation Plan):

# 训练后量化示例 (Post-training quantization example)
def quantize_model(model_path, quantized_model_path, calibration_data_loader, backend="qnnpack"):
    """
    对预训练的TimeMixer模型进行量化
    Quantize a pre-trained TimeMixer model
    
    Args:
        model_path: 预训练模型路径 (path to pretrained model)
        quantized_model_path: 量化后模型保存路径 (path to save quantized model)
        calibration_data_loader: 用于校准的数据加载器 (data loader for calibration)
        backend: 量化后端，可选qnnpack或fbgemm (quantization backend, qnnpack or fbgemm)
    """
    # 设置量化后端
    torch.backends.quantized.engine = backend
    
    # 加载预训练模型
    model = Model(args)
    model.load_state_dict(torch.load(model_path))
    model.eval()
    
    # 准备量化
    model_fp32_prepared = torch.quantization.prepare(model)
    
    # 使用校准数据集
    for x_enc, x_mark_enc, x_dec, x_mark_dec in calibration_data_loader:
        model_fp32_prepared(x_enc, x_mark_enc, x_dec, x_mark_dec)
    
    # 转换为量化模型
    model_int8 = torch.quantization.convert(model_fp32_prepared)
    
    # 保存量化模型
    torch.save(model_int8.state_dict(), quantized_model_path)
    
    return model_int8


# 量化感知训练示例 (Quantization-aware training example)
class QuantizableTimeMixer(nn.Module):
    def __init__(self, configs):
        super().__init__()
        self.model = Model(configs)
        
        # 添加量化存根 (Add quantization stubs)
        self.quant = torch.quantization.QuantStub()
        self.dequant = torch.quantization.DeQuantStub()
        
    def forward(self, x_enc, x_mark_enc, x_dec, x_mark_dec):
        # 量化输入
        x_enc = self.quant(x_enc)
        x_mark_enc = self.quant(x_mark_enc)
        x_dec = self.quant(x_dec)
        x_mark_dec = self.quant(x_mark_dec)
        
        # 模型推理
        outputs = self.model(x_enc, x_mark_enc, x_dec, x_mark_dec)
        
        # 反量化输出
        return self.dequant(outputs)
        
    def fuse_model(self):
        """融合模型中的操作以提高量化性能 (Fuse operations to improve quantization performance)"""
        for m in self.model.modules():
            if hasattr(m, 'fuse_model'):
                m.fuse_model()

使用方法 (Usage):

# 训练后量化使用方法 (Post-training quantization usage)
python quantize.py --model_path ./checkpoints/best_model.pth --backend qnnpack

# 量化感知训练使用方法 (Quantization-aware training usage)
python train.py --task_name long_term_forecast --quantization_aware True --qat_start_epoch 5

预期效果 (Expected Results):

1. 模型大小减小约75%（FP32→INT8）
2. 推理速度提升2-4倍
3. 内存占用减少约70%
4. 性能下降控制在2%以内

特别注意事项 (Special Considerations):

• 量化可能对某些层的精度影响较大，特别是激活函数
• 不同硬件平台对量化的支持程度不同
• 在量化过程中可能需要特别处理自定义操作

[kevinliu2000]

Thanks for your advice. It’s worth clarifying that this feature is not one of our main contributions. If needed, you’re welcome to submit a pull request—we would greatly appreciate your support.