|
- #!/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
- │ ├── MNIST/processed/test.pt
- │ └── MNIST/processed/training.pt
- │ ├── MNIST/raw/train-images-idx3-ubyte
- │ └── MNIST/raw/train-labels-idx1-ubyte
- │ ├── MNIST/raw/t10k-images-idx3-ubyte
- │ └── MNIST/raw/t10k-labels-idx1-ubyte
- ├── train
- │ ├── MNIST/processed/test.pt
- │ └── MNIST/processed/training.pt
- │ ├── MNIST/raw/train-images-idx3-ubyte
- │ └── MNIST/raw/train-labels-idx1-ubyte
- │ ├── MNIST/raw/t10k-images-idx3-ubyte
- │ └── MNIST/raw/t10k-labels-idx1-ubyte
-
-
-
- 示例选用的预训练模型文件为:mnist_epoch1_0.86.pkl
-
-
- 代码会自动放置在/tmp/code目录下。
- 数据集在界面选择后,会自动放置在/tmp/dataset目录下。
- 预训练模型文件在界面选择后,会自动放置在/tmp/pretrainmodel目录下。
- 输出的模型文件也需要放置在/tmp/output目录下,平台会自动下载/tmp/output目录下的文件。
- 如果选用了多数据集,则应在/tmp/dataset后带上数据集名称,比如/tmp/dataset/MnistDataset_torch/train
- '''
-
- 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
-
- 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')
- #The dataset location is placed under /dataset
- parser.add_argument('--traindata', default="/tmp/dataset/train" ,help='path to train dataset')
- parser.add_argument('--testdata', default="/tmp/dataset/test" ,help='path to test dataset')
-
- 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')
- parser.add_argument('--ckpt_url', default="", help='pretrain model path')
- parser.add_argument('--pretrainmodel', default="/tmp/pretrainmodel/mnist_epoch1_0.86.pkl", help='pretrain model path')
-
-
- if __name__ == '__main__':
- args, unknown = parser.parse_known_args()
- # 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=args.traindata, train=True, transform=ToTensor(),download=False)
- test_dataset = mnist.MNIST(root=args.testdata, 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(args.pretrainmodel):
- checkpoint = torch.load(args.pretrainmodel)
- 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))
- #The model output location is placed under /tmp/output
- state = {'model':model.state_dict(), 'optimizer':optimizer.state_dict(), 'epoch':_epoch+1}
- torch.save(state, '/tmp/output/mnist_epoch{}_{:.2f}.pkl'.format(_epoch+1, correct / _sum))
- print('test:')
- print(os.listdir("/tmp/output"))
|