# 引入模块
import torch_npu
from torch_npu.npu import amp  # 导入AMP模块
from torch_npu.contrib import transfer_to_npu  # 使能自动迁移
import time
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
import torchvision

# 初始化运行device
device = torch.device('cuda:0')


# 定义模型网络
class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.net = nn.Sequential(
            # 卷积层
            nn.Conv2d(in_channels=1, out_channels=16,
                      kernel_size=(3, 3),
                      stride=(1, 1),
                      padding=1),
            # 池化层
            nn.MaxPool2d(kernel_size=2),
            # 卷积层
            nn.Conv2d(16, 32, 3, 1, 1),
            # 池化层
            nn.MaxPool2d(2),
            # 将多维输入一维化
            nn.Flatten(),
            nn.Linear(32 * 7 * 7, 16),
            # 激活函数
            nn.ReLU(),
            nn.Linear(16, 10)
        )

    def forward(self, x):
        return self.net(x)


# 下载数据集
train_data = torchvision.datasets.MNIST(
    root='mnist',
    download=True,
    train=True,
    transform=torchvision.transforms.ToTensor()
)

# 定义训练相关参数
batch_size = 64
torch_npu.npu.set_autocast_enabled = "true"
model = CNN().to(device)  # 定义模型
train_dataloader = DataLoader(train_data, batch_size=batch_size)  # 定义DataLoader
loss_func = nn.CrossEntropyLoss().to(device)  # 定义损失函数
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)  # 定义优化器
### 适配定义
scaler = amp.GradScaler()  # 在模型、优化器定义之后，定义GradScaler
###
epochs = 10  # 设置循环次数

# 设置循环
for epoch in range(epochs):
    for imgs, labels in train_dataloader:
        start_time = time.time()  # 记录训练开始时间
        imgs = imgs.to(device)  # 把img数据放到指定NPU上
        labels = labels.to(device)  # 把label数据放到指定NPU上
        with amp.autocast():
            outputs = model(imgs)  # 前向计算
            loss = loss_func(outputs, labels)  # 损失函数计算
        optimizer.zero_grad()
        ### 适配
        # 进行反向传播前后的loss缩放、参数更新
        scaler.scale(loss).backward()  # loss缩放并反向转播
        scaler.step(optimizer)  # 更新参数（自动unscaling）
        scaler.update()  # 基于动态Loss Scale更新loss_scaling系数
        ###
        loss.backward()  # 损失函数反向计算
        optimizer.step()  # 更新优化器

# 定义保存模型
torch.save({
    'epoch': 10,
    'arch': CNN,
    'state_dict': model.state_dict(),
    'optimizer': optimizer.state_dict(),
}, 'checkpoint.pth.tar')

+--------------------------------------------------------------------------------------------------------+
| npu-smi 23.0.0                                   Version: 23.0.0                                       |
+-------------------------------+-----------------+------------------------------------------------------+
| NPU     Name                  | Health          | Power(W)     Temp(C)           Hugepages-Usage(page) |
| Chip    Device                | Bus-Id          | AICore(%)    Memory-Usage(MB)                        |
+===============================+=================+======================================================+
| 104     310                   | OK              | 12.8         45                0     / 969           |
| 0       0                     | 0000:00:0D.0    | 0            592  / 7759                             |
+===============================+=================+======================================================+
+-------------------------------+-----------------+------------------------------------------------------+
| NPU     Chip                  | Process id      | Process name             | Process memory(MB)        |
+===============================+=================+======================================================+
| No running processes found in NPU 104                                                                  |
+===============================+=================+======================================================+

CANN版本：Ascend-cann-toolkit_8.0.RC2.alpha002_linux-x86_64.run
nnrt版本：Ascend-cann-nnrt_8.0.RC2.alpha002_linux-x86_64.run
Driver版本：A300-3010-npu-driver_23.0.0_linux-x86_64.run

python3 te.py
/home/HwHiAiUser/anaconda3/lib/python3.8/site-packages/torch_npu/contrib/transfer_to_npu.py:211: ImportWarning: 
    *************************************************************************************************************
    The torch.Tensor.cuda and torch.nn.Module.cuda are replaced with torch.Tensor.npu and torch.nn.Module.npu now..
    The torch.cuda.DoubleTensor is replaced with torch.npu.FloatTensor cause the double type is not supported now..
    The backend in torch.distributed.init_process_group set to hccl now..
    The torch.cuda.* and torch.cuda.amp.* are replaced with torch.npu.* and torch.npu.amp.* now..
    The device parameters have been replaced with npu in the function below:
    torch.logspace, torch.randint, torch.hann_window, torch.rand, torch.full_like, torch.ones_like, torch.rand_like, torch.randperm, torch.arange, torch.frombuffer, torch.normal, torch._empty_per_channel_affine_quantized, torch.empty_strided, torch.empty_like, torch.scalar_tensor, torch.tril_indices, torch.bartlett_window, torch.ones, torch.sparse_coo_tensor, torch.randn, torch.kaiser_window, torch.tensor, torch.triu_indices, torch.as_tensor, torch.zeros, torch.randint_like, torch.full, torch.eye, torch._sparse_csr_tensor_unsafe, torch.empty, torch._sparse_coo_tensor_unsafe, torch.blackman_window, torch.zeros_like, torch.range, torch.sparse_csr_tensor, torch.randn_like, torch.from_file, torch._cudnn_init_dropout_state, torch._empty_affine_quantized, torch.linspace, torch.hamming_window, torch.empty_quantized, torch._pin_memory, torch.autocast, torch.load, torch.Generator, torch.Tensor.new_empty, torch.Tensor.new_empty_strided, torch.Tensor.new_full, torch.Tensor.new_ones, torch.Tensor.new_tensor, torch.Tensor.new_zeros, torch.Tensor.to, torch.nn.Module.to, torch.nn.Module.to_empty
    *************************************************************************************************************
    
  warnings.warn(msg, ImportWarning)
/home/HwHiAiUser/anaconda3/lib/python3.8/site-packages/torchvision/io/image.py:13: UserWarning: Failed to load image Python extension: 'libc10_cuda.so: cannot open shared object file: No such file or directory'If you don't plan on using image functionality from `torchvision.io`, you can ignore this warning. Otherwise, there might be something wrong with your environment. Did you have `libjpeg` or `libpng` installed before building `torchvision` from source?
  warn(
Traceback (most recent call last):
  File "te.py", line 54, in <module>
    model = CNN().to(device)  # 定义模型
  File "/home/HwHiAiUser/anaconda3/lib/python3.8/site-packages/torch_npu/contrib/transfer_to_npu.py", line 56, in decorated
    return fn(*args, **kwargs)
  File "/home/HwHiAiUser/anaconda3/lib/python3.8/site-packages/torch_npu/utils/module.py", line 60, in to
    self.cast_weight(device)
  File "/home/HwHiAiUser/anaconda3/lib/python3.8/site-packages/torch_npu/utils/module.py", line 119, in cast_weight
    _format_cast(self, current_class)
  File "/home/HwHiAiUser/anaconda3/lib/python3.8/site-packages/torch_npu/utils/module.py", line 86, in _format_cast
    if torch.npu.is_jit_compile_false():
  File "/home/HwHiAiUser/anaconda3/lib/python3.8/site-packages/torch_npu/npu/npu_config.py", line 129, in is_jit_compile_false
    torch_npu.npu._lazy_init()
  File "/home/HwHiAiUser/anaconda3/lib/python3.8/site-packages/torch_npu/npu/__init__.py", line 201, in _lazy_init
    torch_npu._C._npu_init()
RuntimeError: Unsupported soc version: Ascend310
/home/HwHiAiUser/anaconda3/lib/python3.8/tempfile.py:827: ResourceWarning: Implicitly cleaning up <TemporaryDirectory '/tmp/tmphwqh52ea'>
  _warnings.warn(warn_message, ResourceWarning)