1. 模型概述

UNet是一种用于图像分割的卷积神经网络架构，旨在实现快速且精准的分割。它由编码器、瓶颈层和解码器组成：编码器负责下采样以捕捉图像上下文，瓶颈层提取核心特征，解码器则通过上采样生成掩码。此外，其特有的跳跃连接能有效保留空间信息。

2. 准备运行环境

2.1 软件环境

组件	版本
Python	3.10.19
PyTorch	2.1.0
torch_npu	2.1.0.post13
CANN	8.1.RC1

2.2 硬件环境

设备型号	NPU 配置
Atlas 800T A2	单卡

2.3 准备镜像

镜像环境	镜像地址
公网	swr.cn-southwest-2.myhuaweicloud.com/atelier/pytorch_2_1_ascend:pytorch_2.1.0-cann_8.1.rc1-py_3.10-euler_2.10.11-aarch64-snt9b-20250603154214-4e60e43

2.4 启动镜像

IMAGE_ID=swr.cn-southwest-2.myhuaweicloud.com/atelier/pytorch_2_1_ascend:pytorch_2.1.0-cann_8.1.rc1-py_3.10-euler_2.10.11-aarch64-snt9b-20250603154214-4e60e43
CONTAINER_NAME=unet
docker run -u root --privileged \
 --name ${CONTAINER_NAME} \
 --device /dev/davinci0 \
 --device /dev/davinci_manager \
 --device /dev/devmm_svm \
 --device /dev/hisi_hdc \
 -v /usr/local/dcmi:/usr/local/dcmi \
 -v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi \
 -v /usr/local/Ascend/driver:/usr/local/Ascend/driver \
 -v /etc/ascend_install.info:/etc/ascend_install.info \
 -itd ${IMAGE_ID} /bin/bash

3 运行指导

3.1 创建环境

docker exec -it ${CONTAINER_NAME} bash
conda create -n unet --clone PyTorch-2.1.0
conda activate unet

3.2 迁移适配

直接安装deepchem-ascend二进制包，模型已基于适配代码重新编译成二进制包上传pypi

pip install deepchem-ascend==0.0.6

3.3 测试示例

import numpy as np
import deepchem as dc
from deepchem.models.torch_models import UNetModel

x = np.random.randn(5, 3, 32, 32).astype(np.float32)
y = np.random.rand(5, 1, 32, 32).astype(np.float32)
dataset = dc.data.NumpyDataset(x, y)

model = UNetModel(in_channels=3, out_channels=1)
loss = model.fit(dataset, nb_epoch=5)
predictions = model.predict(dataset)
model_device = next(model.model.parameters()).device
print(f"Model device: {model_device}")

assert model_device.type == 'npu', f"Model should be on NPU, but got {model_device}"
print("NPU test passed!")

3.4 运行测试代码

复制上述测试代码保存到test_unet.py

pytest test_unet.py

1. 模型概述

组件

版本

Python

3.10.19

PyTorch

2.1.0

torch_npu

2.1.0.post13

CANN

8.1.RC1

设备型号

NPU 配置

Atlas 800T A2

单卡

镜像环境

镜像地址

公网

swr.cn-southwest-2.myhuaweicloud.com/atelier/pytorch_2_1_ascend:pytorch_2.1.0-cann_8.1.rc1-py_3.10-euler_2.10.11-aarch64-snt9b-20250603154214-4e60e43

2.4 启动镜像

IMAGE_ID=swr.cn-southwest-2.myhuaweicloud.com/atelier/pytorch_2_1_ascend:pytorch_2.1.0-cann_8.1.rc1-py_3.10-euler_2.10.11-aarch64-snt9b-20250603154214-4e60e43
CONTAINER_NAME=unet
docker run -u root --privileged \
 --name ${CONTAINER_NAME} \
 --device /dev/davinci0 \
 --device /dev/davinci_manager \
 --device /dev/devmm_svm \
 --device /dev/hisi_hdc \
 -v /usr/local/dcmi:/usr/local/dcmi \
 -v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi \
 -v /usr/local/Ascend/driver:/usr/local/Ascend/driver \
 -v /etc/ascend_install.info:/etc/ascend_install.info \
 -itd ${IMAGE_ID} /bin/bash

3.3 测试示例

import numpy as np
import deepchem as dc
from deepchem.models.torch_models import UNetModel

x = np.random.randn(5, 3, 32, 32).astype(np.float32)
y = np.random.rand(5, 1, 32, 32).astype(np.float32)
dataset = dc.data.NumpyDataset(x, y)

model = UNetModel(in_channels=3, out_channels=1)
loss = model.fit(dataset, nb_epoch=5)
predictions = model.predict(dataset)
model_device = next(model.model.parameters()).device
print(f"Model device: {model_device}")

assert model_device.type == 'npu', f"Model should be on NPU, but got {model_device}"
print("NPU test passed!")