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

我正在谷歌云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形网深度为4。不使用BatchNorm层（BatchNorm=False）、conv2d_转置（upconv=False）、残差=False和maxpool=True。所以U-Net只包含Conv2D、Conv2D_转置、Dropout、泄漏ReLU、Max池和级联层

---

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

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

如果批处理大小大于128，则会出现内存耗尽错误。你认为这与太小的批次有关吗？你知道“输出融合”的真正含义吗？输出融合是指任何与卷积或matmul融合的元素操作（matmul现在被降为卷积，所以它总是卷积）。请注意，元素级ops是否融合到卷积LHS输入、RHS输入或输出中并不重要——它始终是输出融合。输出融合良好。在您的HLO配置文件中有很多这样的代码表明编译器做得很好。对于内存耗尽错误，您可以使用memory_viewer工具查看填充量，并对其进行优化。感谢您的解释！但为什么“好”融合的利用率如此之差？顺便说一句，我将输入尺寸减少到128x128个图像，因此可以将批量大小增加到1024个。对于输出融合，利用率也提高到约33%的总体利用率。尽管如此，这仍然不如我预期的好。触发器利用率被标准化为TPU矩阵乘法单元的峰值触发器。在输出融合中，有许多低触发器的元素操作，例如广播和常量。这些降低了触发器的平均利用率。如果你看一看没有融合的单个卷积运算的失败率，它们通常都很高。