diff --git a/gpu/train_for_c2net_fail3.py b/gpu/train_for_c2net_fail3.py new file mode 100755 index 0000000..4d2eb07 --- /dev/null +++ b/gpu/train_for_c2net_fail3.py @@ -0,0 +1,73 @@ +''' +在训练环境中,代码会自动放在/tmp/code目录下,上传的数据集会自动放在/tmp/dataset目录下,模型下载路径默认在/tmp/output下,请将模型输出位置指定到/tmp/model, +启智平台界面会提供/tmp/output目录下的文件下载。 +''' + + +from model import Model +import numpy as np +import torch +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 + +# Training settings +parser = argparse.ArgumentParser(description='PyTorch MNIST Example') +#数据集位置放在/tmp/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') + +if __name__ == '__main__': + args = parser.parse_args() + #日志输出 + print('cuda is available:{}'.format(torch.cuda.is_available())) + device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") + 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() + epoch = args.epoch_size + #日志输出 + print('epoch_size is:{}'.format(epoch)) + for _epoch in range(epoch): + 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() + 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)) + #模型输出位置放在/tmp/output下 + torch.save(model, '/tmp/output/mnist_epoch{}_{:.2f}.pkl'.format(_epoch+1, correct / _sum)) + print("----------this is the end--------") + print(a)