Низкая загрузка ТПУ из-за слияния продукции - PullRequest
Новая
       41

Низкая загрузка ТПУ из-за слияния продукции

0 голосов
/ 03 октября 2018

Я обучаю U-Net на Google Cloud TPU.Это работает, но использование очень низкое.

Из-за того, что я не могу загрузить отслеживаемый профиль здесь (?), Скриншот наиболее медленной части здесь: tpu_utilization

Выходной синтез - самая вредная часть.с 58% времени, но только 12% использования.Следующая часть потребления времени (9%) - это «свертка» с использованием 74%.Я не совсем уверен, какие операции мне нужно настроить для лучшего использования выходного слияния?

Ниже приведен мой код для создания U-Net, может быть, внутри есть медленный слой?: 

class UNet:
def create(self, input, start_ch, depth, inc_rate,
           dropout, batchnorm, maxpool, upconv, residual, leaky_relu_alpha):
    with tf.variable_scope('Generator', reuse=tf.AUTO_REUSE):
        o = self._level_block(input, start_ch, depth, inc_rate, dropout, batchnorm, maxpool, upconv, residual,
                              leaky_relu_alpha)
        out_ch = input.shape[3]
        o = tf.layers.conv2d(o, out_ch, 1)
        o = tf.tanh(o)
        return o

def _conv_block(self, m, dim, bn, res, leaky_relu_alpha, do=0):
    n = tf.layers.conv2d(m, dim, 3, padding='same')
    n = tf.nn.leaky_relu(n, alpha=leaky_relu_alpha)
    n = tf.layers.batch_normalization(n) if bn else n
    n = tf.layers.dropout(n, do) if do else n
    n = tf.layers.conv2d(n, dim, 3, padding='same')
    n = tf.nn.leaky_relu(n, alpha=leaky_relu_alpha)
    n = tf.layers.batch_normalization(n)if bn else n
    return tf.concat([m, n], axis=-1) if res else n

def _level_block(self, m, dim, depth, inc, do, bn, mp, up, res, leaky_relu_alpha):
    if depth > 0:
        n = self._conv_block(m, dim, bn, res, leaky_relu_alpha)
        m = tf.layers.max_pooling2d(n, [2, 2], [2, 2]) if mp else tf.layers.conv2d(n, dim, 3, strides=2, padding='same')
        m = self._level_block(m, int(inc * dim), depth - 1, inc, do, bn, mp, up, res, leaky_relu_alpha)
        if up:
            m = tf.image.resize_nearest_neighbor(m, (2*m.shape[1], 2*m.shape[2]))
            m = tf.layers.conv2d(m, dim, 2, padding='same')
            m = tf.nn.leaky_relu(m, alpha=leaky_relu_alpha)
        else:
            m = tf.layers.conv2d_transpose(m, dim, 3, strides=2, padding='same')
            m = tf.nn.leaky_relu(m, alpha=leaky_relu_alpha)

        n = tf.concat([n, m], axis=-1)
        m = self._conv_block(n, dim, bn, res, leaky_relu_alpha)
    else:
        m = self._conv_block(m, dim, bn, res, leaky_relu_alpha, do)
    return m

Мой размер входного пакета - 128. Глубина U-Net - 4. Использование слоев BatchNorm (batchnorm = False), conv2d_transpose (upconv = False), residual = False и maxpool = True.Таким образом, U-Net состоит только из Conv2D, Conv2D_Transpose, Dropout, Leaky ReLU, Max Pooling и сцепления Layers.

Любая идея, что мне нужно настроить, чтобы получить лучшее использование "выходного объединения"?Или хотя бы что влияет на выходной фьюжн?

1 Ответ

0 голосов
/ 05 октября 2018

Я вижу, что у свертки 114 много отступов, 16 * 1 из 128 * 8.Поскольку ваш глобальный размер пакета равен 128, локальный размер пакета для каждого ядра составляет всего 16 (128/8).Можно ли увеличить размер партии вашей модели до 1024?

...