# First Party import smdistributed.modelparallel.tensorflow as smp tf.random.set_seed(1234) tf.config.optimizer.set_jit(True) cfg = { "microbatches": 2, "partitions": 2, "placement_strategy": "spread", "pipeline": "interleaved", "optimize": "memory", } smp.init(cfg) cache_dir = os.path.join(os.path.expanduser("~"), ".keras", "datasets") if not os.path.exists(cache_dir): try: os.mkdir(cache_dir) except OSError as e: if e.errno == errno.EEXIST and os.path.isdir(cache_dir): pass else: raise # Download and load MNIST dataset. (x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data( "MNIST-data-%d" % smp.rank()
import os import errno import horovod.tensorflow as hvd import tensorflow as tf from tensorflow.keras.layers import Conv2D, Dense, Flatten # Rubik: Import TF2.x API import smdistributed.modelparallel.tensorflow as smp tf.random.set_seed(1234) # Rubik: Initialize smp.init() cache_dir = os.path.join(os.path.expanduser("~"), ".keras", "datasets") if not os.path.exists(cache_dir): try: os.mkdir(cache_dir) except OSError as e: if e.errno == errno.EEXIST and os.path.isdir(cache_dir): pass else: raise # Download and load MNIST dataset. (x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data( "MNIST-data-%d" % smp.rank() ) x_train, x_test = x_train / 255.0, x_test / 255.0 # Add a channels dimension
# Standard Library import filecmp import os import shutil import time # First Party import smdistributed.modelparallel.tensorflow as smp import smdistributed.modelparallel.tensorflow.utils as utils smp.init({"partitions": 2, "fork": False}) def generate_big_random_bin_file(filename, size): """ generate big binary file with the specified size in bytes :param filename: the filename :param size: the size in bytes :return:void """ with open("%s" % filename, "wb+") as fout: fout.write(os.urandom(size)) start_time = time.time() src_root_dir = "./send_receive_checkpoint_test" dst_root_dir = "./send_receive_checkpoint_result" filename = "data.bin" if smp.rank() != 0: file_path = os.path.join(src_root_dir, "mp_rank_" + str(smp.rank()))