From e96acc26f703b053a66b6b967a3bc176f8e716d3 Mon Sep 17 00:00:00 2001 From: Robert Nishihara Date: Thu, 25 Jan 2018 14:14:37 -0800 Subject: [PATCH] Fix MNIST downloading problems in parameter server examples. (#1457) * Fix MNIST downloading problems in parameter server examples. * Improve seeding. * Fixes. --- examples/parameter_server/async_parameter_server.py | 10 ++++------ examples/parameter_server/model.py | 12 ++++++++++++ examples/parameter_server/sync_parameter_server.py | 7 ++----- 3 files changed, 18 insertions(+), 11 deletions(-) diff --git a/examples/parameter_server/async_parameter_server.py b/examples/parameter_server/async_parameter_server.py index 56b239a75..3e1aa67da 100644 --- a/examples/parameter_server/async_parameter_server.py +++ b/examples/parameter_server/async_parameter_server.py @@ -3,11 +3,9 @@ from __future__ import division from __future__ import print_function import argparse -from tensorflow.examples.tutorials.mnist import input_data import time import ray - import model parser = argparse.ArgumentParser(description="Run the asynchronous parameter " @@ -35,9 +33,9 @@ class ParameterServer(object): @ray.remote -def worker_task(ps, batch_size=50): +def worker_task(ps, worker_index, batch_size=50): # Download MNIST. - mnist = input_data.read_data_sets("MNIST_data", one_hot=True) + mnist = model.download_mnist_retry(seed=worker_index) # Initialize the model. net = model.SimpleCNN() @@ -65,10 +63,10 @@ if __name__ == "__main__": ps = ParameterServer.remote(all_keys, all_values) # Start some training tasks. - worker_tasks = [worker_task.remote(ps) for _ in range(args.num_workers)] + worker_tasks = [worker_task.remote(ps, i) for i in range(args.num_workers)] # Download MNIST. - mnist = input_data.read_data_sets("MNIST_data", one_hot=True) + mnist = model.download_mnist_retry() i = 0 while True: diff --git a/examples/parameter_server/model.py b/examples/parameter_server/model.py index 7ce1e0033..d68cb4713 100644 --- a/examples/parameter_server/model.py +++ b/examples/parameter_server/model.py @@ -8,6 +8,18 @@ from __future__ import print_function import ray import tensorflow as tf +from tensorflow.examples.tutorials.mnist import input_data +import time + + +def download_mnist_retry(seed=0, max_num_retries=20): + for _ in range(max_num_retries): + try: + return input_data.read_data_sets("MNIST_data", one_hot=True, + seed=seed) + except tf.errors.AlreadyExistsError: + time.sleep(1) + raise Exception("Failed to download MNIST.") class SimpleCNN(object): diff --git a/examples/parameter_server/sync_parameter_server.py b/examples/parameter_server/sync_parameter_server.py index d86c22a22..77c7a4867 100644 --- a/examples/parameter_server/sync_parameter_server.py +++ b/examples/parameter_server/sync_parameter_server.py @@ -3,9 +3,7 @@ from __future__ import division from __future__ import print_function import argparse - import numpy as np -from tensorflow.examples.tutorials.mnist import input_data import ray import model @@ -36,8 +34,7 @@ class Worker(object): def __init__(self, worker_index, batch_size=50): self.worker_index = worker_index self.batch_size = batch_size - self.mnist = input_data.read_data_sets("MNIST_data", one_hot=True, - seed=worker_index) + self.mnist = model.download_mnist_retry(seed=worker_index) self.net = model.SimpleCNN() def compute_gradients(self, weights): @@ -60,7 +57,7 @@ if __name__ == "__main__": for worker_index in range(args.num_workers)] # Download MNIST. - mnist = input_data.read_data_sets("MNIST_data", one_hot=True) + mnist = model.download_mnist_retry() i = 0 current_weights = ps.get_weights.remote()