OpenI_Cloudbrain_Example/gcu_mnist_example/train.py

#!/usr/bin/python
#coding=utf-8
'''
If there are Chinese comments in the code，please add at the beginning：
#!/usr/bin/python
#coding=utf-8   

示例选用的数据集是MnistDataset_torch.zip
数据集结构是：
 MnistDataset_torch.zip
  ├── test
  └── train  

预训练模型文件夹结构是：
Torch_MNIST_Example_Model
├── mnist_epoch1.pkl  

'''

import torch
from model import Model
import numpy as np
from torchvision.datasets import mnist
from torch.nn import CrossEntropyLoss
from torch.optim import SGD
from torch.utils.data import DataLoader
from torchvision.transforms import ToTensor
import argparse
import os
#导入c2net包
from c2net.context import prepare, upload_output

import importlib.util

def is_torch_dtu_available():
    if importlib.util.find_spec("torch_dtu") is None:
        return False
    if importlib.util.find_spec("torch_dtu.core") is None:
        return False
    return importlib.util.find_spec("torch_dtu.core.dtu_model") is not None

# Training settings
parser = argparse.ArgumentParser(description='PyTorch MNIST Example')
parser.add_argument('--epoch_size', type=int, default=1, help='how much epoch to train')
parser.add_argument('--batch_size', type=int, default=256, help='how much batch_size in epoch')

if __name__ == '__main__':
    args, unknown = parser.parse_known_args()

    #初始化导入数据集和预训练模型到容器内
    c2net_context = prepare()
    #获取数据集路径
    MnistDataset_torch_path = c2net_context.dataset_path+"/"+"MnistDataset_torch"
    #获取预训练模型路径
    Torch_MNIST_Example_Model_path = c2net_context.pretrain_model_path+"/"+"GCU_MNIST_Example_Model"
    #获取输出路径
    output_path = c2net_context.output_path
    # load DPU envs-xx.sh
    DTU_FLAG = True
    if is_torch_dtu_available():
        import torch_dtu
        import torch_dtu.distributed as dist
        import torch_dtu.core.dtu_model as dm
        from torch_dtu.nn.parallel import DistributedDataParallel as torchDDP
        print('dtu is available: True') 
        device = dm.dtu_device()
        DTU_FLAG = True
    else:
        print('dtu is available: False')
        device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
        DTU_FLAG = False
        
    
    # 参数声明
    model = Model().to(device)
    optimizer = SGD(model.parameters(), lr=1e-1)
    args, unknown = parser.parse_known_args()
    #log output
    batch_size = args.batch_size
    train_dataset = mnist.MNIST(root=os.path.join(MnistDataset_torch_path, "train"), train=True, transform=ToTensor(),download=False)
    test_dataset = mnist.MNIST(root=os.path.join(MnistDataset_torch_path, "test"), train=False, transform=ToTensor(),download=False)
    train_loader = DataLoader(train_dataset, batch_size=batch_size)
    test_loader = DataLoader(test_dataset, batch_size=batch_size)
    model = Model().to(device)
    sgd = SGD(model.parameters(), lr=1e-1)
    cost = CrossEntropyLoss()
    epochs = args.epoch_size
    print('epoch_size is:{}'.format(epochs))

    # 如果有保存的模型，则加载模型，并在其基础上继续训练
    if os.path.exists(os.path.join(Torch_MNIST_Example_Model_path, "mnist_epoch1_0.81.pkl")):
        checkpoint = torch.load(os.path.join(Torch_MNIST_Example_Model_path, "mnist_epoch1_0.81.pkl"))
        model.load_state_dict(checkpoint['model'])
        optimizer.load_state_dict(checkpoint['optimizer'])
        start_epoch = checkpoint['epoch']
        print('加载 epoch {} 权重成功！'.format(start_epoch))
    else:
        start_epoch = 0
        print('无保存模型，将从头开始训练！')

    for _epoch in range(start_epoch, epochs):
        print('the {} epoch_size begin'.format(_epoch + 1))
        model.train()
        for idx, (train_x, train_label) in enumerate(train_loader):
            train_x = train_x.to(device)
            train_label = train_label.to(device)
            label_np = np.zeros((train_label.shape[0], 10))
            sgd.zero_grad()
            predict_y = model(train_x.float())
            loss = cost(predict_y, train_label.long())
            if idx % 10 == 0:
                print('idx: {}, loss: {}'.format(idx, loss.sum().item()))
            loss.backward()
            if DTU_FLAG:
                dm.optimizer_step(sgd, barrier=True)
            else:
                sgd.step()
            
            
        correct = 0
        _sum = 0
        model.eval()
        for idx, (test_x, test_label) in enumerate(test_loader):
            test_x = test_x
            test_label = test_label
            predict_y = model(test_x.to(device).float()).detach()
            predict_ys = np.argmax(predict_y.cpu(), axis=-1)
            label_np = test_label.numpy()
            _ = predict_ys == test_label
            correct += np.sum(_.numpy(), axis=-1)
            _sum += _.shape[0]
        print('accuracy: {:.2f}'.format(correct / _sum))
        state = {'model':model.state_dict(), 'optimizer':optimizer.state_dict(), 'epoch':_epoch+1}
        torch.save(state, '{}/mnist_epoch{}_{:.2f}.pkl'.format(output_path, _epoch+1, correct / _sum))
        print(os.listdir('{}'.format(output_path)))