Распараллеливание кратных независимых моделей в tenorflow-gpu / keras

Question

Мне нужно обучить ряд моделей, но я не получаю выгоды от ускорения графического процессора, используя tenorflow-gpu / keras, поскольку время линейно увеличивается с количеством обучаемых моделей.

В

class Models(tf.keras.Model):

    def __init__(self,N_MODELS=1):
        super(Models, self).__init__()

        self.block_i = [estimate_affine() 
                        for node in range(N_MODELS)]


    def call(self, inputs):

        x = [self.block_i[i](input_i) for i,input_i in enumerate(inputs)]

        return x

список слоев N_MODELS составлен и, как идентификаторы, должен быть распараллелен.Поскольку это не так, даже несмотря на то, что результат - это то, что я ожидаю, я думаю, что моя реализация не оптимальна.Есть идеи, как сделать его распараллеливаемым?

Best

Paul

Вот тойнет N_MODELS линейной регрессии

import tensorflow as tf
tf.enable_eager_execution()
from tensorflow.keras import layers
import numpy as np
from numpy import random
import time


class estimate_affine(layers.Layer):
    def __init__(self):
        '''
        '''
        super(estimate_affine, self).__init__()

        self.a = tf.Variable(initial_value=[0.], dtype='float32',trainable=True,name='par1')
        self.b = tf.Variable(initial_value=[0.], dtype='float32',trainable=True,name='par2')


    def call(self, inputs):
        return (self.a,self.b)

class Models(tf.keras.Model):

    def __init__(self,N_MODELS=1):
        super(Models, self).__init__()

        self.block_i = [estimate_affine() 
                        for node in range(N_MODELS)]


    def call(self, inputs):

        x = [self.block_i[i](input_i) for i,input_i in enumerate(inputs)]

        return x


N_ITERATIONS=100
N_POINTS=100
ls_t=[]



for N_MODELS in [5,10,50,100,1000]:
    t=time.time()


    ### Aim is to fit N_MODELS on  N_POINTS which are basically N_MODELS of ax+b 
    a=np.random.randint(0,10,N_MODELS)
    b=np.random.randint(0,10,N_MODELS)
    noise=np.random.rand(N_POINTS) * 1
    x=np.linspace(0,1,N_POINTS)
    dataset=np.array([a_i  *( x + noise) + b_i    for a_i,b_i in zip(a,b)])

    model=Models(N_MODELS=N_MODELS)
    optimizer=tf.keras.optimizers.SGD(learning_rate=5e-3)

    for i in range(N_ITERATIONS):
        with tf.GradientTape() as tape:
            outputs=model(dataset)
            L=tf.reduce_sum([((outputs[idx][0]*x+outputs[idx][1]) 
                                   - dataset[idx,:])**2 for idx in range(N_MODELS)])

        grads = tape.gradient(L, model.trainable_variables)
        optimizer.apply_gradients(zip(grads, model.trainable_variables))
    t_diff=time.time()-t
    print('N_MODEL : {}, time : {}'.format(N_MODELS,t_diff))

    ls_t.append(t_diff)