multi gpu 실패, tensorflow mirrored strategy

멀티 GPU가 잡히지 않을 떄

mirrored_strategy = tf.distribute.MirroredStrategy()

tf.keras로 multi gpu를 쓰려면 위에 구문을 쳐주고

mirrored_strategy scope내에 model 생성 및 compile을 해줘야 한다.

with mirrored_strategy.scope():
    dense = densenet.DenseNet121(include_top=False, classes=2)
    inputs = tf.keras.Input(shape=(256, 256, 3))
    x = dense(inputs)
    x = tf.keras.layers.GlobalAveragePooling2D()(x)
    outputs = tf.keras.layers.Dense(1, activation=tf.nn.sigmoid)(x)
    model = tf.keras.Model(inputs=inputs, outputs=outputs)
    
    model.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy'])

그런데 가끔, 아래와 같이 gpu를 사용 안 하는 오류가 있다.

WARNING:tensorflow:There are non-GPU devices in `tf.distribute.Strategy`, not using nccl allreduce.

그럴땐 아래 구문으로 mirrored_strategy를 생성해주면 된다.

GPU가 잡힐 때

mirrored_strategy = tf.distribute.MirroredStrategy(devices=["/gpu:0","/gpu:1","/gpu:2","/gpu:3"],
                      cross_device_ops=tf.distribute.HierarchicalCopyAllReduce())