[tune/sgd] Fix Jenkins (#5765)

ci/jenkins_tests/run_tune_tests.sh

@@ -63,12 +63,10 @@ $SUPPRESS_OUTPUT docker run --rm --shm-size=${SHM_SIZE} --memory=${MEMORY_SIZE}
     --smoke-test
 
 $SUPPRESS_OUTPUT docker run --rm --shm-size=${SHM_SIZE} --memory=${MEMORY_SIZE} $DOCKER_SHA \
-    pip install tensorflow==2.0.0rc1 && python /ray/python/ray/tune/examples/async_hyperband_example.py \
-    --smoke-test
+    bash -c 'pip install tensorflow==2.0.0rc1 && python /ray/python/ray/tune/examples/async_hyperband_example.py --smoke-test'
 
 $SUPPRESS_OUTPUT docker run --rm --shm-size=${SHM_SIZE} --memory=${MEMORY_SIZE} $DOCKER_SHA \
-    pip install tensorflow==1.15.0rc1 && python /ray/python/ray/tune/examples/async_hyperband_example.py \
-    --smoke-test
+    bash -c 'pip install tensorflow==1.15.0rc1 && python /ray/python/ray/tune/examples/async_hyperband_example.py --smoke-test'
 
 $SUPPRESS_OUTPUT docker run --rm --shm-size=${SHM_SIZE} --memory=${MEMORY_SIZE} $DOCKER_SHA \
     python /ray/python/ray/tune/examples/tune_mnist_ray_hyperband.py \

docker/examples/Dockerfile

@@ -7,11 +7,12 @@ RUN conda install -y numpy
 RUN apt-get install -y zlib1g-dev
 # The following is needed to support TensorFlow 1.14
 RUN conda remove -y --force wrapt
+RUN pip install -U pip
 RUN pip install gym[atari] opencv-python-headless tensorflow lz4 keras pytest-timeout smart_open tensorflow_probability
 RUN pip install -U h5py  # Mutes FutureWarnings
 RUN pip install --upgrade bayesian-optimization
 RUN pip install --upgrade git+git://github.com/hyperopt/hyperopt.git
-RUN pip install --upgrade sigopt nevergrad scikit-optimize hpbandster ConfigSpace lightgbm xgboost
-RUN pip install -U mlflow
+RUN pip install --upgrade sigopt nevergrad scikit-optimize hpbandster ConfigSpace lightgbm xgboost torch torchvision
+RUN pip install -U tabulate mlflow
 RUN pip install -U pytest-remotedata>=0.3.1
 RUN conda install pytorch-cpu torchvision-cpu -c pytorch

docker/tune_test/Dockerfile

@@ -5,18 +5,18 @@ FROM ray-project/base-deps
 # We install ray and boto3 to enable the ray autoscaler as
 # a test runner.
 RUN conda install -y numpy
+RUN pip install -U pip
 RUN pip install -U https://ray-wheels.s3-us-west-2.amazonaws.com/latest/ray-0.8.0.dev5-cp36-cp36m-manylinux1_x86_64.whl boto3
 # We install this after the latest wheels -- this should not override the latest wheels.
 RUN apt-get install -y zlib1g-dev
 # The following is needed to support TensorFlow 1.14
 RUN conda remove -y --force wrapt
-RUN pip install gym[atari]==0.10.11 opencv-python-headless tensorflow lz4 keras pytest-timeout smart_open
+RUN pip install gym[atari]==0.10.11 opencv-python-headless tensorflow lz4 keras pytest-timeout smart_open torch torchvision
 RUN pip install --upgrade bayesian-optimization
 RUN pip install --upgrade git+git://github.com/hyperopt/hyperopt.git
 RUN pip install --upgrade sigopt nevergrad scikit-optimize hpbandster ConfigSpace lightgbm xgboost
 RUN pip install -U mlflow
 RUN pip install -U pytest-remotedata>=0.3.1
-RUN conda install pytorch-cpu torchvision-cpu -c pytorch
 
 # RUN mkdir -p /root/.ssh/
 

python/ray/experimental/sgd/examples/cifar_tf_example.py

@@ -63,7 +63,7 @@ def create_model(config):
     model.add(Dropout(0.25))
 
     model.add(Flatten())
-    model.add(Dense(512))
+    model.add(Dense(64))
     model.add(Activation("relu"))
     model.add(Dropout(0.5))
     model.add(Dense(num_classes))
@@ -160,10 +160,7 @@ if __name__ == "__main__":
         default=1,
         help="Sets number of replicas for training.")
     parser.add_argument(
-        "--batch-size",
-        type=int,
-        default=512,
-        help="Sets number of replicas for training.")
+        "--batch-size", type=int, default=32, help="Sets batch size.")
     parser.add_argument(
         "--use-gpu",
         action="store_true",