Question

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

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

输出融合是最有害的部分。 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和Concatenation层组成。

任何想法我需要调整以获得更好的“输出融合”利用率吗？或者至少是什么会影响输出融合？

Answer 1

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

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

1 个答案: