train.py

import tensorlayerx as tlx
from tensorlayerx.dataflow import DataLoader
from tensorlayerx.optimizers.lr import LRScheduler
from tensorlayerx.vision.transforms import Compose
from tqdm import tqdm

from tlxzoo.datasets import CocoHumanPoseEstimationDataset
from tlxzoo.module.hrnet import *
from tlxzoo.vision.human_pose_estimation import HumanPoseEstimation


pck = PCK()


def valid(model, test_data):
    model.set_eval()
    accs = []
    for idx, batch in enumerate(tqdm(test_data)):
        images = batch[0]
        target, target_weight = batch[1]["target"]
        y_pred = model(images)
        _, avg_accuracy, _, _ = pck(network_output=y_pred, target=target)
        if avg_accuracy == 0:
            continue
        accs.append(avg_accuracy)

    print((sum(accs) * 1.0) / len(accs))


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 batch in train_dataset:
                network.set_train()

                with tf.GradientTape() as tape:
                    # compute outputs
                    _logits = network(batch['image'])
                    _loss_ce = loss_fn(_logits, target=batch['target'], target_weight=batch['target_weight'])

                grad = tape.gradient(_loss_ce, train_weights)

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

                _, avg_accuracy, _, _ = pck(network_output=_logits, target=batch['target'])
                train_acc += avg_accuracy

                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))
                    print("   train acc: {}".format(train_acc / n_iter))
                    print("   learning rate: ", optimizer.lr().numpy())

            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))
                print("   train acc: {}".format(train_acc / n_iter))

            optimizer.lr.step()


class EpochDecay(LRScheduler):
    def __init__(self, learning_rate, last_epoch=0, verbose=False):
        super(EpochDecay, self).__init__(learning_rate, last_epoch, verbose)

    def get_lr(self):

        if int(self.last_epoch) >= 65:
            return self.base_lr * 0.01

        if int(self.last_epoch) >= 40:
            return self.base_lr * 0.1

        return self.base_lr


if __name__ == '__main__':
    transform = Compose([
        Gather(),
        Crop(),
        Resize((256, 256)),
        Normalize(),
        GenerateTarget()
    ])
    train_dataset = CocoHumanPoseEstimationDataset(root='./data/coco2017', split='train', transform=transform)
    train_dataloader = DataLoader(train_dataset, batch_size=2)
    test_dataset = CocoHumanPoseEstimationDataset(root='./data/coco2017', split='test', transform=transform)
    test_dataloader = DataLoader(test_dataset, batch_size=2)

    model = HumanPoseEstimation("hrnet")

    scheduler = EpochDecay(1e-3)
    optimizer = tlx.optimizers.Adam(lr=scheduler)
    # optimizer = tlx.optimizers.SGD(lr=scheduler)

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

    model.save_weights("./demo/vision/human_pose_estimation/hrnet/model.npz")