<jupyter_start><jupyter_text>ExperimentThis is only the experiment about using distribute strategy.    Using MirroredStrategy eagerly has significant overhead currently.   SO DON'T RUN ON MULTI-GPU NOW.   When using @tf.function to execute in a graph mode, this code can't use Tensor.numpy() function.    That means functions in `detection/utils/misc.py` don't work.   And for now, I can't find an alternative.<jupyter_code>import os
import tensorflow as tf
import numpy as np
import visualize

# tensorflow config - using one gpu and extending the GPU 
# memory region needed by the TensorFlow process
os.environ['CUDA_VISIBLE_DEVICES'] = '0,1'
# config = tf.ConfigProto()
# config.gpu_options.allow_growth = True
# session = tf.Session(config=config)<jupyter_output><empty_output><jupyter_text>load dataset<jupyter_code>from detection.datasets import coco, data_generator
img_mean = (123.675, 116.28, 103.53)
# img_std = (58.395, 57.12, 57.375)
img_std = (1., 1., 1.)
train_dataset = coco.CocoDataSet('./COCO2017/', 'val',
                                 flip_ratio=0.5,
                                 pad_mode='fixed',
                                 mean=img_mean,
                                 std=img_std,
                                 scale=(640, 832),
                                 debug=True)

train_generator = data_generator.DataGenerator(train_dataset)<jupyter_output>loading annotations into memory...
Done (t=0.55s)
creating index...
index created!<jupyter_text>load model<jupyter_code>from detection.models.detectors import faster_rcnn

strategy = tf.distribute.MirroredStrategy()

with strategy.scope():
    
    model = faster_rcnn.FasterRCNN(
        num_classes=len(train_dataset.get_categories()))
    
    img, img_meta, _, _ = train_dataset[0]
    batch_imgs = tf.Variable(np.expand_dims(img, 0))
    batch_metas = tf.Variable(np.expand_dims(img_meta, 0))
    _ = model((batch_imgs, batch_metas), training=False)
    
    model.load_weights('weights/faster_rcnn.h5', by_name=True)<jupyter_output>INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Redu[...]<jupyter_text>use tf.data<jupyter_code>global_batch_size = 2

with strategy.scope():
    train_tf_dataset = tf.data.Dataset.from_generator(
        train_generator, (tf.float32, tf.float32, tf.float32, tf.int32))
    train_tf_dataset = train_tf_dataset.padded_batch(
        global_batch_size, padded_shapes=([None, None, None], [None], [None, None], [None]))
    train_tf_dataset = train_tf_dataset.prefetch(100).shuffle(100)
    
    dist_dataset = strategy.experimental_distribute_dataset(train_tf_dataset)<jupyter_output><empty_output><jupyter_text>train model<jupyter_code>with strategy.scope():

    def train_step(inputs):
        batch_imgs, batch_metas, batch_bboxes, batch_labels = inputs
        with tf.GradientTape() as tape:
            rpn_class_loss, rpn_bbox_loss, rcnn_class_loss, rcnn_bbox_loss = \
                model((batch_imgs, batch_metas, batch_bboxes, batch_labels), training=True)

            loss_value = rpn_class_loss + rpn_bbox_loss + rcnn_class_loss + rcnn_bbox_loss
            loss = tf.reduce_sum(loss_value) * (1.0 / global_batch_size)


        grads = tape.gradient(loss_value, model.trainable_variables)
        optimizer.apply_gradients(zip(grads, model.trainable_variables))
        
        return loss_value
    
with strategy.scope():
 
    def distributed_train_step(inputs):
        per_replica_losses = strategy.experimental_run_v2(train_step,
                                                          args=(inputs,))
        return strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses,
                               axis=None)
    
    
    optimizer = tf.keras.optimizers.SGD(1e-3, momentum=0.9, nesterov=True)

    epochs = 12

    for epoch in range(epochs):

        loss_history = []
        for (batch, inputs) in enumerate(train_tf_dataset):
            
            batch_loss = distributed_train_step(inputs)

            loss_history.append(batch_loss.numpy())

            print('epoch:', epoch, ' batch:', batch, ' loss:', np.mean(loss_history))
with strategy.scope():
    from detection.datasets.utils import get_original_image
    ori_img = get_original_image(img, img_meta, img_mean)
    proposals = model.simple_test_rpn(img, img_meta)
    res = model.simple_test_bboxes(img, img_meta, proposals)
    visualize.display_instances(ori_img, res['rois'], res['class_ids'], 
                            train_dataset.get_categories(), scores=res['scores'])<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.11683348498635122, 'fraction_numerical': 0.02183803457688808, 'mean_word_length': 3.56857855361596, 'pattern_counts': {'":': 0, '<': 19, '<?xml version=': 0, '>': 19, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 8, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '5', 'dir': 'jupyter-structured-clean-dedup', 'id': '590108', 'license': 'MIT', 'n_tokens_mistral': 1869, 'n_tokens_neox': 1806, 'n_words': 320, 'path': 'train_model_multigpu.ipynb', 'repo_name': 'pskurochkin/tf-eager-fasterrcnn', 'size': '5495'}