diff --git a/python/ray/tune/examples/mnist_pytorch.py b/python/ray/tune/examples/mnist_pytorch.py index a5fe48e5d1ba..a1cb983b112c 100644 --- a/python/ray/tune/examples/mnist_pytorch.py +++ b/python/ray/tune/examples/mnist_pytorch.py @@ -8,7 +8,6 @@ import torch.nn.functional as F import torch.optim as optim from torchvision import datasets, transforms -from torch.autograd import Variable # Training settings parser = argparse.ArgumentParser(description='PyTorch MNIST Example') @@ -120,7 +119,6 @@ def train(epoch): for batch_idx, (data, target) in enumerate(train_loader): if args.cuda: data, target = data.cuda(), target.cuda() - data, target = Variable(data), Variable(target) optimizer.zero_grad() output = model(data) loss = F.nll_loss(output, target) @@ -131,16 +129,17 @@ def test(): model.eval() test_loss = 0 correct = 0 - for data, target in test_loader: - if args.cuda: - data, target = data.cuda(), target.cuda() - data, target = Variable(data, volatile=True), Variable(target) - output = model(data) - test_loss += F.nll_loss( - output, target, size_average=False).item() # sum up batch loss - pred = output.data.max( - 1, keepdim=True)[1] # get the index of the max log-probability - correct += pred.eq(target.data.view_as(pred)).long().cpu().sum() + with torch.no_grad(): + for data, target in test_loader: + if args.cuda: + data, target = data.cuda(), target.cuda() + output = model(data) + # sum up batch loss + test_loss += F.nll_loss(output, target, reduction='sum').item() + # get the index of the max log-probability + pred = output.argmax(dim=1, keepdim=True) + correct += pred.eq( + target.data.view_as(pred)).long().cpu().sum() test_loss = test_loss / len(test_loader.dataset) accuracy = correct.item() / len(test_loader.dataset) @@ -176,7 +175,8 @@ def test(): "training_iteration": 1 if args.smoke_test else 20 }, "resources_per_trial": { - "cpu": 3 + "cpu": 3, + "gpu": int(not args.no_cuda) }, "run": "train_mnist", "num_samples": 1 if args.smoke_test else 10, diff --git a/python/ray/tune/examples/mnist_pytorch_trainable.py b/python/ray/tune/examples/mnist_pytorch_trainable.py index 24fc4951dd37..52aed6596abf 100644 --- a/python/ray/tune/examples/mnist_pytorch_trainable.py +++ b/python/ray/tune/examples/mnist_pytorch_trainable.py @@ -9,7 +9,6 @@ import torch.nn.functional as F import torch.optim as optim from torchvision import datasets, transforms -from torch.autograd import Variable from ray.tune import Trainable @@ -127,7 +126,6 @@ def _train_iteration(self): for batch_idx, (data, target) in enumerate(self.train_loader): if self.args.cuda: data, target = data.cuda(), target.cuda() - data, target = Variable(data), Variable(target) self.optimizer.zero_grad() output = self.model(data) loss = F.nll_loss(output, target) @@ -138,18 +136,17 @@ def _test(self): self.model.eval() test_loss = 0 correct = 0 - for data, target in self.test_loader: - if self.args.cuda: - data, target = data.cuda(), target.cuda() - data, target = Variable(data, volatile=True), Variable(target) - output = self.model(data) - - # sum up batch loss - test_loss += F.nll_loss(output, target, size_average=False).item() - - # get the index of the max log-probability - pred = output.data.max(1, keepdim=True)[1] - correct += pred.eq(target.data.view_as(pred)).long().cpu().sum() + with torch.no_grad(): + for data, target in self.test_loader: + if self.args.cuda: + data, target = data.cuda(), target.cuda() + output = self.model(data) + # sum up batch loss + test_loss += F.nll_loss(output, target, reduction='sum').item() + # get the index of the max log-probability + pred = output.argmax(dim=1, keepdim=True) + correct += pred.eq( + target.data.view_as(pred)).long().cpu().sum() test_loss = test_loss / len(self.test_loader.dataset) accuracy = correct.item() / len(self.test_loader.dataset) @@ -188,7 +185,8 @@ def _restore(self, checkpoint_path): "training_iteration": 1 if args.smoke_test else 20, }, "resources_per_trial": { - "cpu": 3 + "cpu": 3, + "gpu": int(not args.no_cuda) }, "run": TrainMNIST, "num_samples": 1 if args.smoke_test else 20, diff --git a/test/jenkins_tests/run_multi_node_tests.sh b/test/jenkins_tests/run_multi_node_tests.sh index ed7e1c8bc95c..07221fbb70fd 100755 --- a/test/jenkins_tests/run_multi_node_tests.sh +++ b/test/jenkins_tests/run_multi_node_tests.sh @@ -363,12 +363,11 @@ docker run --rm --shm-size=${SHM_SIZE} --memory=${MEMORY_SIZE} $DOCKER_SHA \ --smoke-test docker run --rm --shm-size=${SHM_SIZE} --memory=${MEMORY_SIZE} $DOCKER_SHA \ - python /ray/python/ray/tune/examples/mnist_pytorch.py \ - --smoke-test + python /ray/python/ray/tune/examples/mnist_pytorch.py --smoke-test --no-cuda docker run --rm --shm-size=${SHM_SIZE} --memory=${MEMORY_SIZE} $DOCKER_SHA \ python /ray/python/ray/tune/examples/mnist_pytorch_trainable.py \ - --smoke-test + --smoke-test --no-cuda docker run --rm --shm-size=${SHM_SIZE} --memory=${MEMORY_SIZE} $DOCKER_SHA \ python /ray/python/ray/tune/examples/genetic_example.py \