Python 带有自定义训练步骤的Tensorflow分布式训练

我面临着缓慢的训练运行，我试图通过使用Tensorflow的

策略

API来利用所有4个GPU来扩大训练过程

我正在使用

MirrorStrategy

和

experimental\u distribute\u dataset

对数据集进行分区

我的训练数据的本质是稀疏矩阵和密集矩阵的混合。（从数据中选择随机指数）。然而，在我当前版本的TF（2.1）中。

稀疏_矩阵

没有静态大小，是一个

参差不齐的

张量

这一点很难看，也是一种解决方法，但我将我的

稀疏矩阵列表

直接传递给

训练

函数，并通过将随机索引推到

生成器

中来建立全局队列索引

现在这个方法工作得很好，但是太慢了，我想尝试使用所有GPU进行训练。由于我必须手动将

sparse\u matrix\u列表

划分为

num\u workers

splits，因此问题更加严重

但是，目前的主要问题是，培训过程似乎不是并行的，副本（GPU）似乎是按顺序运行的。。我通过

nvidia smi

验证了这一点，并登录了

train_进程

功能

我以前没有分布式培训的经验，也不知道为什么会出现这种情况，如果有人能找到更好的方法来处理

空闲

和

密集

数据的混合，我会非常感激。我目前在获取数据方面面临巨大瓶颈，GPU利用率不足（在10-30%之间波动）

def distributed_train_step(inputs, sparse_matrix_list):
    per_replica_losses = strategy.experimental_run_v2(train_process, args=(inputs, sparse_matrix_list)
    return strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses,
                                axis=None)


def train_process(inputs, sparse_matrix_list):
    worker_id = tf.distribute.get_replica_context().replica_id_in_sync_group
    replica_batch_size = inputs.shape[0]
    slice_start = replica_batch_size * worker_id
    replica_sparse_matrix = sparse_matrix_list[slice_start:slice_start + replica_batch_size]

    return train_step(inputs, replica_sparse_matrix)


def train_step(inputs, sparse_matrix_list):
     with tf.GradientTape() as tape:
           outputs, mu, sigma, feat_out, logit = model(inputs)
           loss = K.backend.mean(custom_loss(inputs, sparse_matrix_list)

     return loss

def get_batch_data(sparse_matrix_list):
    # Queue with the random indices into the training data (List of Lists with each 
    # entry len == batch_size)
    # train_indicie is a global q
    next_batch_indicies = train_indicies.get()

    batch_sparse_list = sparse_matrix_list[next_batch_indicies]
            

dist_dataset = strategy.experimental_distribute_dataset(train_dataset)

for batch, inputs in enumerate(dist_dataset, 1):
    # sparse_matrix_list is passed to this main "train" function from outside this module.
    batch_sparse_matrix_slice = get_batch_data(sparse_matrix_list)

    loss = distributed_train_step(inputs, batch_sparse_matrix_slice)