From d430e901ca64e55bbdea489d5fde6398d3df7e77 Mon Sep 17 00:00:00 2001 From: wjtest1201 Date: Tue, 30 Aug 2022 17:54:01 +0800 Subject: [PATCH] =?UTF-8?q?=E6=9B=B4=E6=96=B0=20'gpu/train.py'?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- gpu/train.py | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/gpu/train.py b/gpu/train.py index c615864..f4c09c9 100755 --- a/gpu/train.py +++ b/gpu/train.py @@ -39,10 +39,14 @@ parser.add_argument('--testdata', default="/dataset/test" ,help='path to test da 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') +def gettime(): + timestr = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S') + return timestr + if __name__ == '__main__': args, unknown = parser.parse_known_args() #log output - print('cuda is available:{}'.format(torch.cuda.is_available())) + print(gettime(), '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) @@ -53,9 +57,9 @@ if __name__ == '__main__': sgd = SGD(model.parameters(), lr=1e-1) cost = CrossEntropyLoss() epoch = args.epoch_size - print('epoch_size is:{}'.format(epoch)) + print(gettime(), 'epoch_size is:{}'.format(epoch)) for _epoch in range(epoch): - print('the {} epoch_size begin'.format(_epoch + 1)) + print(gettime(), '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) @@ -66,6 +70,7 @@ if __name__ == '__main__': loss = cost(predict_y, train_label.long()) #if idx % 10 == 0: #print('idx: {}, loss: {}'.format(idx, loss.sum().item())) + print(gettime()) print('idx: {}, loss: {}'.format(idx, loss.sum().item())) loss.backward() sgd.step() @@ -82,6 +87,6 @@ if __name__ == '__main__': _ = predict_ys == test_label correct += np.sum(_.numpy(), axis=-1) _sum += _.shape[0] - print('accuracy: {:.2f}'.format(correct / _sum)) + print(gettime(), 'accuracy: {:.2f}'.format(correct / _sum)) #The model output location is placed under /model torch.save(model, '/model/mnist_epoch{}_{:.2f}.pkl'.format(_epoch+1, correct / _sum))