由于输出融合，TPU 利用率低

Question

我正在 Google Cloud TPU 上训练 U-Net。可以，但是利用率很低。

由于我无法在此处上传跟踪的配置文件（？），最慢部分的屏幕截图在这里：

输出融合是最有害的部分。 58% 的时间，但只有 12% 的利用率。下一个耗时部分（9%）是“卷积”，利用率为 74%。我不确定需要调整哪些操作才能更好地利用输出融合？

下面是我创建U-Net的代码，可能里面有一个slow layer？：

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 和 Concatenation 层。

知道我需要调整什么以获得更好的“输出融合”利用率吗？或者至少是什么影响了输出融合？

Answer 1

我可以看到卷积。114 有很多填充，128*8 中有 16*1。由于您的全局批处理大小为 128，因此每个内核的本地批处理大小仅为 16 (128/8)。您是否可以将模型的批量大小增加到 1024？