UniTS 昇腾适配完成总结

1. 核心适配代码

• ✅ utils/ascend_utils.py: 昇腾设备工具模块
• ✅ utils/device_utils.py: 设备兼容性工具（统一 CUDA/昇腾/CPU 接口）
• ✅ utils/memory_utils.py: 内存工具模块（设备兼容性内存操作）
• ✅ models/ascend_adapter.py: 模型适配层

2. 核心文件适配

• ✅ run.py: 主运行脚本已适配
• ✅ run_pretrain.py: 预训练脚本已适配
• ✅ exp/exp_sup.py: 监督学习实验类已适配
• ✅ exp/exp_pretrain.py: 预训练实验类已适配

3. 主要修改点

设备选择

• 替换 torch.cuda.device_count() → get_device_count()
• 替换 torch.cuda.current_device() → get_current_device()
• 替换 .to(device_id) → .to(device) (使用设备对象而非 ID)

设备操作

• 替换 torch.cuda.synchronize() → synchronize()
• 替换 torch.cuda.empty_cache() → empty_cache()
• 替换 torch.cuda.memory_reserved() → memory_reserved()

分布式训练

• 适配 DDP (DistributedDataParallel) 以支持昇腾设备
• 保持与 CUDA 的兼容性

4. 算子适配

• UniTS.py

  • 文件头 from utils.ascend_utils import scaled_dot_product_attention_compat

  • 将所有 F.scaled_dot_product_attention 改成 scaled_dot_product_attention_compat

📋 适配统计

代码修改统计

• 设备选择相关: ~10 处
• CUDA 调用替换: ~20 处
• 设备对象使用: ~30 处
• 内存操作: ~5 处

🔧 关键适配点

1. 设备检测和选择

# 旧代码
device_id = dist.get_rank() % torch.cuda.device_count()
model = model.to(device_id)

# 新代码（适配后）
device_count = get_device_count()
device_id = dist.get_rank() % device_count
device = get_device_by_id(device_id)
model = model.to(device)

2. 设备操作统一

# 旧代码
torch.cuda.synchronize()
torch.cuda.empty_cache()

# 新代码（适配后）
synchronize()
empty_cache()

3. 分布式训练适配

# 适配后的 DDP
if str(self.device).startswith('ascend'):
    model = nn.parallel.DistributedDataParallel(
        model, find_unused_parameters=True)
else:
    model = nn.parallel.DistributedDataParallel(
        model, device_ids=[self.device_id], find_unused_parameters=True)

使用方法

NPU调试指南

本指南以 Supervised learning 任务为例，执行命令是：

bash ./scripts/supervised_learning/UniTS_supervised.sh

环境

1. 系统环境

   EulerOS 2.1(SP10) CANN-8.0.0

2. python环境

   pip install -r requirements.txt

   保证 torch 与 torch_npu 版本一致

2. 数据准备

1. 打开 download_data_all.sh ，按需求复制文档中的数据url进行手动下载（需要科学上网），本指南以 https://www.timeseriesclassification.com/aeon-toolkit/Archives/Multivariate2018_ts.zip 为例。

2. 下载完成后，解压缩，将数据中的文件夹拷贝到 UniTS-npu/dataset/ 路径下。以下以 JapaneseVowels 为例。

3. 数据配置

1. 将 multi_task.yaml 拷贝一份，重命名为： example.yaml。

2. 打开 example.yaml，文件包含多个数据集信息，往下拉（315行附近）找到并保留 JapaneseVowels 相关的，将其他删掉。确保配置中的 root_path 是正确的。

4 修改启动脚本

1. 对于 supervised_learning 任务，打开 ./scripts/supervised_learning/UniTS_supervised.sh

   • 设置 wandb_mode=disabled

   • 设置 --task_data_config_path 的正确路径：data_provider/example.yaml

执行命令进行验证

bash ./scripts/supervised_learning/UniTS_supervised.sh

若执行成功，则其他任务可参考以上方法修改对应的文件以达到目标。

总结

1. 对代码而言，从GPU到NPU，将涉及到GPU的函数在NPU上进行同名替换，遇到特别的或当前算子不支持的函数功能，则需要重新实现一个可支持的。

2. 根据不同任务，修改的主要文件是：data_provider/ 下的 yaml 文件以及对应任务的启动脚本 .sh 文件。查看该路径下的data_loader.py可了解更多关于数据加载的信息。

3. 确保 .sh 中配置文件路径的准确性，以及 yaml 文件中数据路径的准确性。

4. 如果硬件环境只有一张卡，则在 .sh 中，要设置 --nproc-per-node 1 .