<commit_before>#
# Copyright 2018 Analytics Zoo Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

# Reference: https://mxnet.apache.org/api/python/docs/tutorials/packages/gluon/image/mnist.html

import argparse

from zoo import init_spark_on_local, init_spark_on_yarn
from zoo.ray import RayContext
from zoo.orca.learn.mxnet import MXNetTrainer, create_trainer_config


def get_data_iters(config, kv):
    import os
    import zipfile
    import mxnet as mx
    from bigdl.dataset.base import maybe_download

    # In order to avoid conflict where multiple workers on the same node download and
    # zip data under the same location, here we let each worker have its own folder.

    # Not using mxnet.test_utils.get_mnist_iterator directly because data path is
    # hard-coded in this function.

    # In practice, data is supposed to be stored on a file system accessible to workers on
    # all nodes, for example, on HDFS or S3.
    maybe_download("mnist.zip", "worker" + str(kv.rank),
                   "http://data.mxnet.io/mxnet/data/mnist.zip")
    if not os.path.isdir("worker" + str(kv.rank) + "/data"):
        with zipfile.ZipFile("worker" + str(kv.rank) + "/mnist.zip") as zf:
            zf.extractall("worker" + str(kv.rank) + "/data")

    train_iter = mx.io.MNISTIter(
        image="worker" + str(kv.rank) + "/data/train-images-idx3-ubyte",
        label="worker" + str(kv.rank) + "/data/train-labels-idx1-ubyte",
        input_shape=(1, 28, 28),
        batch_size=config["batch_size"],
        shuffle=True,
        flat=False,
        num_parts=kv.num_workers,
        part_index=kv.rank)
    val_iter = mx.io.MNISTIter(
        image="worker" + str(kv.rank) + "/data/t10k-images-idx3-ubyte",
        label="worker" + str(kv.rank) + "/data/t10k-labels-idx1-ubyte",
        input_shape=(1, 28, 28),
        batch_size=config["batch_size"],
        flat=False,
        num_parts=kv.num_workers,
        part_index=kv.rank)
    return train_iter, val_iter


def get_model(config):
    import mxnet as mx
    from mxnet import gluon
    from mxnet.gluon import nn
    import mxnet.ndarray as F

    class LeNet(gluon.Block):
        def __init__(self, **kwargs):
            super(LeNet, self).__init__(**kwargs)
            with self.name_scope():
                # layers created in name_scope will inherit name space
                # from parent layer.
                self.conv1 = nn.Conv2D(20, kernel_size=(5, 5))
                self.pool1 = nn.MaxPool2D(pool_size=(2, 2), strides=(2, 2))
                self.conv2 = nn.Conv2D(50, kernel_size=(5, 5))
                self.pool2 = nn.MaxPool2D(pool_size=(2, 2), strides=(2, 2))
                self.fc1 = nn.Dense(500)<commit_msg>Use filelock in MXNet LeNet example (#2372)<commit_after>#
# Copyright 2018 Analytics Zoo Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

# Reference: https://mxnet.apache.org/api/python/docs/tutorials/packages/gluon/image/mnist.html

import argparse

from zoo import init_spark_on_local, init_spark_on_yarn
from zoo.ray import RayContext
from zoo.orca.learn.mxnet import MXNetTrainer, create_trainer_config


def get_data_iters(config, kv):
    from mxnet.test_utils import get_mnist_iterator
    from filelock import FileLock
    with FileLock("data.lock"):
        return get_mnist_iterator(config["batch_size"], (1, 28, 28),
                                  num_parts=kv.num_workers, part_index=kv.rank)


def get_model(config):
    import mxnet as mx
    from mxnet import gluon
    from mxnet.gluon import nn
    import mxnet.ndarray as F

    class LeNet(gluon.Block):
        def __init__(self, **kwargs):
            super(LeNet, self).__init__(**kwargs)
            with self.name_scope():
                # layers created in name_scope will inherit name space
                # from parent layer.
                self.conv1 = nn.Conv2D(20, kernel_size=(5, 5))
                self.pool1 = nn.MaxPool2D(pool_size=(2, 2), strides=(2, 2))
                self.conv2 = nn.Conv2D(50, kernel_size=(5, 5))
                self.pool2 = nn.MaxPool2D(pool_size=(2, 2), strides=(2, 2))
                self.fc1 = nn.Dense(500)
{'fraction_non_alphanumeric': 0.10306677018633541, 'fraction_numerical': 0.019409937888198756, 'mean_word_length': 2.874436090225564, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 2, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5208661', 'n_tokens_mistral': 1681, 'n_tokens_neox': 1591, 'n_words': 477}