# 引入模块 

import time 

import torch_npu 

import torch 

import torch.nn as nn 

from torch.utils.data import Dataset, DataLoader 

import torchvision 

 

# 初始化运行device 

device = torch.device('npu: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    

model = CNN().to(f'npu:{device}' if isinstance(device, int) else device)  # 定义模型 

train_dataloader = DataLoader(train_data, batch_size=batch_size)    # 定义DataLoader 

loss_func = nn.CrossEntropyLoss().to(f'npu:{device}' if isinstance(device, int) else device)    # 定义损失函数 

optimizer = torch.optim.SGD(model.parameters(), lr=0.1)    # 定义优化器 

epochs = 10  # 设置循环次数 

start_time = time.time()  

# 设置循环 

for epoch in range(epochs): 

    for i,(imgs, labels) in  enumerate(train_dataloader): 

        #start_time = time.time()    # 记录训练开始时间 

        imgs = imgs.to(f'npu:{device}' if isinstance(device, int) else device)    # 把img数据放到指定NPU上 

        labels = labels.to(f'npu:{device}' if isinstance(device, int) else device)    # 把label数据放到指定NPU上 

        outputs = model(imgs)    # 前向计算 

        loss = loss_func(outputs, labels)    # 损失函数计算 

        optimizer.zero_grad() 

        loss.backward()    # 损失函数反向计算 

        optimizer.step()    # 更新优化器 

    print(f"Epoch [{epoch+1}/{epochs}], Step [{i+1}/{len(train_dataloader)}], Loss: {loss.item():.4f}") 

 

# 定义保存模型 

torch.save({ 

               'epoch': 10, 

               'arch': CNN, 

               'state_dict': model.state_dict(), 

               'optimizer' : optimizer.state_dict(), 

            },'checkpoint.pth.tar') 

end_time = time.time()   

elapsed_time = end_time - start_time  # 计算耗时 

print(f"该函数执行耗时: {elapsed_time}秒") 

////////错误提示

Optype [MaxPoolWithArgmaxV1] of Ops kernel [AIcoreEngine] is unsupported