Почему у меня растёт время работы, когда я тренирую простой RNN на GPU, а время работы линейно увеличивается на CPU?

Question

Вот небольшой код, который вы сможете запустить на своем компьютере. Я не интересуюсь тем, что нейронная сеть (NN) изучает сейчас, только во время ее работы. (Под временем выполнения я подразумеваю время, затрачиваемое оптимизатором (прямое + обратное распространение) в одном отдельном учебном примере).

import tensorflow as tf
import numpy as np
import time

#%% Parameters :
size_image = 64
nb_unroll  = 10

#%% Auxiliary function:
def create_radom_images(size_image):
    a, b, r = size_image//2, size_image//2, size_image//4
    y,x = np.ogrid[-a:size_image-a,-b:size_image-b]
    mask = x*x + y*y <= r*r    
    im_x = np.zeros((size_image, size_image));im_x[mask] = 1
    im_x = im_x[np.newaxis,...,np.newaxis]
    im_y = im_x + np.random.rand(1,size_image,size_image,1)
    return im_x, im_y

tf.compat.v1.reset_default_graph()

# Create Placeholders of the correct shape
print('Creation of the placeholders X, Y and learning_rate')
X = tf.compat.v1.placeholder(tf.float64, shape=(1, size_image, size_image, 1), name = 'X')
Y = tf.compat.v1.placeholder(tf.float64, shape=(1, size_image, size_image, 1), name = 'Y_true')

#%% nb_unroll wrapped
Y_hat = tf.ones_like(X)
for i in range(1,nb_unroll+1):
     print('unrolled block: '+ str(i) + ' -- iteration: ' + str(1))
     # Block i
     # convolutions with learnt bowties
     size_fiter = size_image * 2
     filter_tensor = tf.compat.v1.get_variable('my_filter_'+str(i),\
                         shape = [size_fiter,size_fiter,1,1],\
                         dtype = tf.float64,\
                         initializer = tf.contrib.layers.xavier_initializer(seed = 0),\
                         trainable = True)
     Y_hat = tf.add(X,tf.nn.conv2d(Y_hat, filter_tensor, strides = [1,1,1,1], padding = 'SAME'))

cost = tf.compat.v1.losses.mean_squared_error(Y,Y_hat)

# Backpropagation: Define the tensorflow optimizer. Use an AdamOptimizer that minimizes the cost.
print('Definition of the optimiser')
opt_adam = tf.compat.v1.train.AdamOptimizer(learning_rate=0.001) #GradientDescentOptimizer
update_ops = tf.compat.v1.get_collection(tf.compat.v1.GraphKeys.UPDATE_OPS)
with tf.control_dependencies(update_ops):
    optimizer = opt_adam.minimize(cost)

init = tf.compat.v1.global_variables_initializer()
with tf.compat.v1.Session() as sess:

    # Run the initialization
    print('Initialisation of the variables')
    sess.run(init)

    print('Loading data')
    minibatch_X, minibatch_Y = create_radom_images(size_image)

    print('Optimising')
    t1 = time.time()                        
    _ , temp_cost = sess.run([optimizer,cost],\
        feed_dict = {X:minibatch_X, Y:minibatch_Y})
    print('Minimisation processed for cost method : ones or tfw')
    t2 = time.time() - t1    
    print('Running session: ' + str(t2) + 's elapsed.')

Важной переменной здесь является nb_unroll. Когда nb_unroll = 1, время работы на GPU невелико (1,45 сек, если size_image = 64), и оно взлетает, как только nb_unroll> 1. Например, я получаю время работы 16,04 сек, если nb_unroll = 2 и size_image = 64. На моем графическом процессоре для изображения 64x64 я в основном получаю следующее время выполнения (в зависимости от nb_unroll):

nb_unroll = 1 -> время выполнения = 1,45 с

nb_unroll = 2 --> время работы = 16,04 с

nb_unroll = 3 -> время работы = 16,06 с

nb_unroll = 4 -> время работы = 16,12 с

nb_unroll = 5-> время бега = 16,15 с

nb_unroll = 10 -> время бега = 16,38 с

Мне не удается объяснить этот скачок с 1,45 с до 16,04 с, а также тот факт, чтодля nb_unroll> = 2 время выполнения почти одинаково, но оно должно увеличиться более значительно. Я мог бы заметить, что когда я устанавливаю size_image на более высокие значения, поведение времени выполнения одинаково: для nb_unroll = 1 время выполнения составляет около нескольких секунд, и как только nb_unroll = 2, время выполнения легко достигает нескольких минут илидаже часы. Но, опять же, время выполнения для nb_unroll = 2,3,4 и т. Д. Более или менее одинаково (немного увеличивается, но незначительно по сравнению с предыдущим прыжком).

Наоборот,когда я запускаю один и тот же код на процессоре, время выполнения линейно увеличивается с nb_unroll.

nb_unroll = 1 -> время работы = 0,86 с

nb_unroll = 2 -> время работы = 1,70 с

nb_unroll = 3 -> время работы = 2,52с

nb_unroll = 4 -> время работы = 3,45 с

nb_unroll = 5 -> время работы = 4,37 с

nb_unroll = 10 -> время работы =8,75 с

Есть ли способ обойти это стремительное время работы с использованием графического процессора?

Большое спасибо за вашу помощь.

PS: Обратите внимание, что у меня естьиспользовать фильтры, которые в два раза больше изображения. Я знаю, что это не обычно в рамках глубокого обучения, но мне нужны фильтры этой формы из математической теории моей модели.