train.py

import tensorlayerx as tlx
from tensorlayerx.dataflow import DataLoader
from tensorlayerx.vision.transforms import Compose

from tlxzoo.datasets import CocoDetectionDataset
from tlxzoo.datasets.coco import CocoEvaluator
from tlxzoo.module.detr import *
from tlxzoo.vision.object_detection import ObjectDetection


def valid(model, data_loader):
    from tqdm import tqdm
    coco_evaluator = CocoEvaluator(data_loader.dataset.coco, "bbox")
    model.set_eval()

    for idx, batch in enumerate(tqdm(data_loader)):
        inputs = batch[0]["inputs"]
        pixel_mask = batch[0]["pixel_mask"]
        labels = batch[1]
        # forward pass
        outputs = model(inputs=inputs)

        orig_target_sizes = tlx.convert_to_tensor([target["orig_size"] for target in batch[1]], dtype=tlx.float32)

        results = post_process(outputs["pred_logits"], outputs["pred_boxes"], orig_target_sizes)

        res = {target['image_id']: output for target, output in zip(labels, results)}
        coco_evaluator.update(res)

    coco_evaluator.synchronize_between_processes()
    coco_evaluator.accumulate()
    coco_evaluator.summarize()


class Trainer(tlx.model.Model):

    def tf_train(
            self, n_epoch, train_dataset, network, loss_fn, train_weights, optimizer, metrics, print_train_batch,
            print_freq, test_dataset
    ):
        import tensorflow as tf
        import time
        for epoch in range(n_epoch):
            start_time = time.time()

            train_loss, train_acc, n_iter = 0, 0, 0
            for X_batch, y_batch in train_dataset:
                network.set_train()

                with tf.GradientTape() as tape:
                    # compute outputs
                    _logits = network(**X_batch)
                    _loss_ce = loss_fn(_logits, y_batch)

                grad = tape.gradient(_loss_ce, train_weights)

                optimizer.apply_gradients(zip(grad, train_weights))
                train_loss += _loss_ce
                n_iter += 1

                if print_train_batch:
                    print("Epoch {} of {} {} took {}".format(epoch + 1, n_epoch, n_iter, time.time() - start_time))
                    print("   train loss: {}".format(train_loss / n_iter))

            if epoch + 1 == 1 or (epoch + 1) % print_freq == 0:
                print("Epoch {} of {} took {}".format(epoch + 1, n_epoch, time.time() - start_time))
                print("   train loss: {}".format(train_loss / n_iter))


if __name__ == '__main__':
    transform = Compose([
        LabelFormatConvert(),
        Resize(size=800, max_size=1333),
        Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225))
    ])
    train_dataset = CocoDetectionDataset(root='./data/coco2017', split='train', transform=transform)
    train_dataloader = DataLoader(train_dataset, batch_size=2, collate_fn=collate_fn)
    test_dataset = CocoDetectionDataset(root='./data/coco2017', split='test', transform=transform)
    test_dataloader = DataLoader(test_dataset, batch_size=2, collate_fn=collate_fn)

    model = ObjectDetection(backbone="detr")

    optimizer = tlx.optimizers.Adam(lr=1e-6)
    metric = None

    trainer = Trainer(network=model, loss_fn=model.loss_fn, optimizer=optimizer, metrics=metric)
    trainer.train(n_epoch=1, train_dataset=train_dataloader, test_dataset=test_dataloader, print_freq=1,
                  print_train_batch=True)

    valid(model, test_dataloader)