混合精度适配网络样例

以下样例为混合精度适配前网络样例，用户需根据具体场景选择适配方法，具体适配操作请参考（可选）混合精度适配。

#引入模块 
import time 
import torch 
import torch.nn as nn 
import torch_npu 
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() 
) 
# 定义batchsize 
batch_size = 64  
# 定义DataLoader 
train_dataloader = DataLoader(train_data, batch_size = batch_size)  

# 把模型放到指定NPU上 
model = CNN().to(device) 
# 定义损失函数  
loss_func = nn.CrossEntropyLoss().to(device) 
# 定义优化器    
optimizer = torch.optim.SGD(model.parameters(), lr = 0.1) 

# 设置循环 
epochs = 10 
for i 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上 
        outputs = model(imgs)    # 前向计算 
        loss = loss_func(outputs, labels)    # 损失函数计算 
        optimizer.zero_grad() 
        loss.backward()          # 损失函数反向计算 
        optimizer.step()         # 更新优化器

PATH = "state_dict_model.pt"        # 创建保存路径 
torch.save(model.state_dict(), PATH)# 保存模型