def main(epochs, enable_function, buffer_size, batch_size, download_path, num_examples=70000, embedding_dim=256, enc_units=1024, dec_units=1024): file_path = utils.download(download_path) train_ds, test_ds, inp_lang, targ_lang = utils.create_dataset( file_path, num_examples, buffer_size, batch_size) vocab_inp_size = len(inp_lang.word_index) + 1 vocab_tar_size = len(targ_lang.word_index) + 1 encoder = nmt.Encoder(vocab_inp_size, embedding_dim, enc_units, batch_size) decoder = nmt.Decoder(vocab_tar_size, embedding_dim, dec_units) train_obj = Train(epochs, enable_function, encoder, decoder, inp_lang, targ_lang, batch_size, batch_size) print ('Training ...') return train_obj.training_loop(train_ds, test_ds)
def main(epochs, enable_function, buffer_size, batch_size, download_path, num_examples=70000, embedding_dim=256, enc_units=1024, dec_units=1024, num_gpu=1): devices = ['/device:GPU:{}'.format(i) for i in range(num_gpu)] strategy = tf.distribute.MirroredStrategy(devices) num_replicas = strategy.num_replicas_in_sync with strategy.scope(): file_path = utils.download(download_path) train_ds, test_ds, inp_lang, targ_lang = utils.create_dataset( file_path, num_examples, buffer_size, batch_size) vocab_inp_size = len(inp_lang.word_index) + 1 vocab_tar_size = len(targ_lang.word_index) + 1 num_train_steps_per_epoch = tf.data.experimental.cardinality(train_ds) num_test_steps_per_epoch = tf.data.experimental.cardinality(test_ds) train_iterator = strategy.make_dataset_iterator(train_ds) test_iterator = strategy.make_dataset_iterator(test_ds) local_batch_size, remainder = divmod(batch_size, num_replicas) template = ('Batch size ({}) must be divisible by the ' 'number of replicas ({})') if remainder: raise ValueError(template.format(batch_size, num_replicas)) encoder = nmt.Encoder(vocab_inp_size, embedding_dim, enc_units, local_batch_size) decoder = nmt.Decoder(vocab_tar_size, embedding_dim, dec_units) train_obj = DistributedTrain(epochs, enable_function, encoder, decoder, inp_lang, targ_lang, batch_size, local_batch_size) print('Training ...') return train_obj.training_loop(train_iterator, test_iterator, num_train_steps_per_epoch, num_test_steps_per_epoch, strategy)
def main(epochs, enable_function, buffer_size, batch_size, download_path, num_examples=70000, embedding_dim=256, enc_units=1024, dec_units=1024, num_gpu=1): devices = ['/device:GPU:{}'.format(i) for i in range(num_gpu)] strategy = tf.distribute.MirroredStrategy(devices) num_replicas = strategy.num_replicas_in_sync with strategy.scope(): file_path = utils.download(download_path) train_ds, test_ds, inp_lang, targ_lang = utils.create_dataset( file_path, num_examples, buffer_size, batch_size) vocab_inp_size = len(inp_lang.word_index) + 1 vocab_tar_size = len(targ_lang.word_index) + 1 num_train_steps_per_epoch = tf.data.experimental.cardinality(train_ds) num_test_steps_per_epoch = tf.data.experimental.cardinality(test_ds) train_iterator = strategy.make_dataset_iterator(train_ds) test_iterator = strategy.make_dataset_iterator(test_ds) local_batch_size, remainder = divmod(batch_size, num_replicas) template = ('Batch size ({}) must be divisible by the ' 'number of replicas ({})') if remainder: raise ValueError(template.format(batch_size, num_replicas)) encoder = nmt.Encoder(vocab_inp_size, embedding_dim, enc_units, local_batch_size) decoder = nmt.Decoder(vocab_tar_size, embedding_dim, dec_units) train_obj = DistributedTrain(epochs, enable_function, encoder, decoder, inp_lang, targ_lang, batch_size, local_batch_size) print ('Training ...') return train_obj.training_loop(train_iterator, test_iterator, num_train_steps_per_epoch, num_test_steps_per_epoch, strategy)