I'm trying to implement a ResNeXt block (type B) from the paper in TensorFlow:

To do so, I've used a loop to generate each path depending on the cardinality parameter (cardinality = 32 in the image above)

def resnext_block(input_tensor, filters, stride=1, projection=False, cardinality=32, depth=4):
    residual = input_tensor

    # Split
    branches = []
    for i in range(cardinality):
        with tf.variable_scope("branch%i" % i):
            branch = tf.layers.conv2d(inputs=input_tensor, filters=depth, kernel_size=1, padding='same')
            branch = tf.layers.batch_normalization(branch, training=train_flag)
            branch = activation(branch)
            branch = tf.layers.conv2d(inputs=branch, filters=depth, kernel_size=3, strides=stride, padding='same')
            branch = tf.layers.batch_normalization(branch, training=train_flag)
            branch = activation(branch)
            branches.append(branch)

    # Merge
    merged = tf.concat(branches, -1)
    # Transform
    transformed = tf.layers.conv2d(inputs=merged, filters=filters, kernel_size=1, padding='same')
    transformed = tf.layers.batch_normalization(transformed, training=train_flag)

    if projection:
        residual = tf.layers.conv2d(inputs=input_tensor, filters=filters, kernel_size=1, strides=stride)
        residual = tf.layers.batch_normalization(residual, training=train_flag)

    return activation(transformed + residual)

If I run my model for low values of cardinality (< 4), the graph builds quickly and training gets underway.

For high values of cardinality, TensorFlow spends a very long time building the graph (i.e. with 1 CPU core at 100% and no other activity). The events file grows very large (> 1GB) presumably because there are so many nodes in the graph (despite the fact the number of trainable parameters can be kept to a similar number by reducing depth).

Are there any tricks I'm missing to speed up/optimise this process?

I'm not seeing very many ResNeXt examples in TensorFlow (the ones I have found have cardinality set to small numbers), which makes me think maybe the graph architecture in TF is not suited for these types of models.

Update

It looks like grouped convolutions are the best way to achieve this (block style C in the paper). However these are not yet implemented in TF 1.10. See this GitHub issue