ram_train_eval.py

import os
import time
from datetime import timedelta
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

from mmengine.config import Config
from mmengine.utils import ProgressBar
from transformers import AutoConfig, AutoModel

class RamDataset(torch.utils.data.Dataset):
    def __init__(self, data_path, is_train=True, num_relation_classes=56):
        super().__init__()
        self.num_relation_classes = num_relation_classes
        data = np.load(data_path, allow_pickle=True)
        self.samples = data["arr_0"]
        sample_num = self.samples.size
        self.sample_idx_list = []
        for idx in range(sample_num):
            if self.samples[idx]["is_train"] == is_train:
                self.sample_idx_list.append(idx)

    def __getitem__(self, idx):
        sample = self.samples[self.sample_idx_list[idx]]
        object_num = sample["feat"].shape[0]
        embedding = torch.from_numpy(sample["feat"])
        gt_rels = sample["relations"]
        rel_target = self._get_target(object_num, gt_rels)
        return embedding, rel_target, gt_rels

    def __len__(self):
        return len(self.sample_idx_list)

    def _get_target(self, object_num, gt_rels):
        rel_target = torch.zeros([self.num_relation_classes, object_num, object_num])
        for ii, jj, cls_relationship in gt_rels:
            rel_target[cls_relationship, ii, jj] = 1
        return rel_target


class RamModel(nn.Module):
    def __init__(
        self,
        pretrained_model_name_or_path,
        load_pretrained_weights=True,
        num_transformer_layer=2,
        input_feature_size=256,
        output_feature_size=768,
        cls_feature_size=512,
        num_relation_classes=56,
        pred_type="attention",
        loss_type="bce",
    ):
        super().__init__()
        # 0. config
        self.cls_feature_size = cls_feature_size
        self.num_relation_classes = num_relation_classes
        self.pred_type = pred_type
        self.loss_type = loss_type

        # 1. fc input and output
        self.fc_input = nn.Sequential(
            nn.Linear(input_feature_size, output_feature_size),
            nn.LayerNorm(output_feature_size),
        )
        self.fc_output = nn.Sequential(
            nn.Linear(output_feature_size, output_feature_size),
            nn.LayerNorm(output_feature_size),
        )
        # 2. transformer model
        if load_pretrained_weights:
            self.model = AutoModel.from_pretrained(pretrained_model_name_or_path)
        else:
            config = AutoConfig.from_pretrained(pretrained_model_name_or_path)
            self.model = AutoModel.from_config(config)
        if num_transformer_layer != "all" and isinstance(num_transformer_layer, int):
            self.model.encoder.layer = self.model.encoder.layer[:num_transformer_layer]
        # 3. predict head
        self.cls_sub = nn.Linear(output_feature_size, cls_feature_size * num_relation_classes)
        self.cls_obj = nn.Linear(output_feature_size, cls_feature_size * num_relation_classes)
        # 4. loss
        if self.loss_type == "bce":
            self.bce_loss = nn.BCEWithLogitsLoss()
        elif self.loss_type == "multi_label_ce":
            print("Use Multi Label Cross Entropy Loss.")

    def forward(self, embeds, attention_mask=None):
        """
        embeds: (batch_size, token_num, feature_size)
        attention_mask: (batch_size, token_num)
        """
        # 1. fc input
        embeds = self.fc_input(embeds)
        # 2. transformer model
        position_ids = torch.ones([1, embeds.shape[1]]).to(embeds.device).to(torch.long)
        outputs = self.model.forward(inputs_embeds=embeds, attention_mask=attention_mask, position_ids=position_ids)
        embeds = outputs["last_hidden_state"]
        # 3. fc output
        embeds = self.fc_output(embeds)
        # 4. predict head
        batch_size, token_num, feature_size = embeds.shape
        sub_embeds = self.cls_sub(embeds).reshape([batch_size, token_num, self.num_relation_classes, self.cls_feature_size]).permute([0, 2, 1, 3])
        obj_embeds = self.cls_obj(embeds).reshape([batch_size, token_num, self.num_relation_classes, self.cls_feature_size]).permute([0, 2, 1, 3])
        if self.pred_type == "attention":
            cls_pred = sub_embeds @ torch.transpose(obj_embeds, 2, 3) / self.cls_feature_size**0.5  # noqa
        elif self.pred_type == "einsum":
            cls_pred = torch.einsum("nrsc,nroc->nrso", sub_embeds, obj_embeds)
        return cls_pred

    def loss(self, pred, target, attention_mask):
        loss_dict = dict()
        batch_size, relation_num, _, _ = pred.shape

        mask = torch.zeros_like(pred).to(pred.device)
        for idx in range(batch_size):
            n = torch.sum(attention_mask[idx]).to(torch.int)
            mask[idx, :, :n, :n] = 1
        pred = pred * mask - 9999 * (1 - mask)

        if self.loss_type == "bce":
            loss = self.bce_loss(pred, target)
        elif self.loss_type == "multi_label_ce":
            input_tensor = torch.permute(pred, (1, 0, 2, 3))
            target_tensor = torch.permute(target, (1, 0, 2, 3))
            input_tensor = pred.reshape([relation_num, -1])
            target_tensor = target.reshape([relation_num, -1])
            loss = self.multilabel_categorical_crossentropy(target_tensor, input_tensor)
            weight = loss / loss.max()
            loss = loss * weight
        loss = loss.mean()
        loss_dict["loss"] = loss

        # running metric
        recall_20 = get_recall_N(pred, target, object_num=20)
        loss_dict["recall@20"] = recall_20
        return loss_dict

    def multilabel_categorical_crossentropy(self, y_true, y_pred):
        """
        https://kexue.fm/archives/7359
        """
        y_pred = (1 - 2 * y_true) * y_pred
        y_pred_neg = y_pred - y_true * 9999
        y_pred_pos = y_pred - (1 - y_true) * 9999
        zeros = torch.zeros_like(y_pred[..., :1])
        y_pred_neg = torch.cat([y_pred_neg, zeros], dim=-1)
        y_pred_pos = torch.cat([y_pred_pos, zeros], dim=-1)
        neg_loss = torch.logsumexp(y_pred_neg, dim=-1)
        pos_loss = torch.logsumexp(y_pred_pos, dim=-1)
        return neg_loss + pos_loss


def get_recall_N(y_pred, y_true, object_num=20):
    """
    y_pred: [batch_size, 56, object_num, object_num]
    y_true: [batch_size, 56, object_num, object_num]
    """

    device = y_pred.device
    recall_list = []

    for idx in range(len(y_true)):
        sample_y_true = []
        sample_y_pred = []

        # find topk
        _, topk_indices = torch.topk(
            y_true[idx : idx + 1].reshape(
                [
                    -1,
                ]
            ),
            k=object_num,
        )
        for index in topk_indices:
            pred_cls = index // (y_true.shape[2] ** 2)
            index_subject_object = index % (y_true.shape[2] ** 2)
            pred_subject = index_subject_object // y_true.shape[2]
            pred_object = index_subject_object % y_true.shape[2]
            if y_true[idx, pred_cls, pred_subject, pred_object] == 0:
                continue
            sample_y_true.append([pred_subject, pred_object, pred_cls])

        # find topk
        _, topk_indices = torch.topk(
            y_pred[idx : idx + 1].reshape(
                [
                    -1,
                ]
            ),
            k=object_num,
        )
        for index in topk_indices:
            pred_cls = index // (y_pred.shape[2] ** 2)
            index_subject_object = index % (y_pred.shape[2] ** 2)
            pred_subject = index_subject_object // y_pred.shape[2]
            pred_object = index_subject_object % y_pred.shape[2]
            sample_y_pred.append([pred_subject, pred_object, pred_cls])

        recall = len([x for x in sample_y_pred if x in sample_y_true]) / (len(sample_y_true) + 1e-8)
        recall_list.append(recall)

    recall = torch.tensor(recall_list).to(device).mean() * 100
    return recall


class RamTrainer(object):
    def __init__(self, config):
        self.config = config
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self._build_dataset()
        self._build_dataloader()
        self._build_model()
        self._build_optimizer()
        self._build_lr_scheduler()

    def _build_dataset(self):
        self.dataset = RamDataset(**self.config.dataset)

    def _build_dataloader(self):
        self.dataloader = torch.utils.data.DataLoader(
            self.dataset,
            batch_size=self.config.dataloader.batch_size,
            shuffle=True if self.config.dataset.is_train else False,
        )

    def _build_model(self):
        self.model = RamModel(**self.config.model).to(self.device)
        if self.config.load_from is not None:
            self.model.load_state_dict(torch.load(self.config.load_from))
        self.model.train()

    def _build_optimizer(self):
        self.optimizer = torch.optim.AdamW(self.model.parameters(), lr=self.config.optim.lr, weight_decay=self.config.optim.weight_decay, eps=self.config.optim.eps, betas=self.config.optim.betas)

    def _build_lr_scheduler(self):
        self.lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(self.optimizer, milestones=self.config.optim.lr_scheduler.step, gamma=self.config.optim.lr_scheduler.gamma)

    def train(self):
        t_start = time.time()
        running_avg_loss = 0
        for epoch_idx in range(self.config.num_epoch):
            for batch_idx, batch_data in enumerate(self.dataloader):
                batch_embeds = batch_data[0].to(torch.float32).to(self.device)
                batch_target = batch_data[1].to(torch.float32).to(self.device)
                attention_mask = batch_embeds.new_ones((batch_embeds.shape[0], batch_embeds.shape[1]))
                batch_pred = self.model.forward(batch_embeds, attention_mask)
                loss_dict = self.model.loss(batch_pred, batch_target, attention_mask)
                loss = loss_dict["loss"]
                recall_20 = loss_dict["recall@20"]
                self.optimizer.zero_grad()
                loss.backward()
                torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.config.optim.max_norm, self.config.optim.norm_type)
                self.optimizer.step()
                running_avg_loss += loss.item()

                if batch_idx % 100 == 0:
                    t_current = time.time()
                    num_finished_step = epoch_idx * self.config.num_epoch * len(self.dataloader) + batch_idx + 1
                    num_to_do_step = (self.config.num_epoch - epoch_idx - 1) * len(self.dataloader) + (len(self.dataloader) - batch_idx - 1)
                    avg_speed = num_finished_step / (t_current - t_start)
                    eta = num_to_do_step / avg_speed
                    print(
                        "ETA={:0>8}, Epoch={}, Batch={}/{}, LR={}, Loss={:.4f}, RunningAvgLoss={:.4f}, Recall@20={:.2f}%".format(
                            str(timedelta(seconds=int(eta))), epoch_idx + 1, batch_idx, len(self.dataloader), self.lr_scheduler.get_last_lr()[0], loss.item(), running_avg_loss / num_finished_step, recall_20.item()
                        )
                    )
            self.lr_scheduler.step()
            if not os.path.exists(self.config.output_dir):
                os.makedirs(self.config.output_dir)
            save_path = os.path.join(self.config.output_dir, "epoch_{}.pth".format(epoch_idx + 1))
            print("Save epoch={} checkpoint to {}".format(epoch_idx + 1, save_path))
            torch.save(self.model.state_dict(), save_path)
        return save_path


class RamPredictor(object):
    def __init__(self, config):
        self.config = config
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self._build_dataset()
        self._build_dataloader()
        self._build_model()

    def _build_dataset(self):
        self.dataset = RamDataset(**self.config.dataset)

    def _build_dataloader(self):
        self.dataloader = torch.utils.data.DataLoader(self.dataset, batch_size=self.config.dataloader.batch_size, shuffle=False)

    def _build_model(self):
        self.model = RamModel(**self.config.model).to(self.device)
        if self.config.load_from is not None:
            self.model.load_state_dict(torch.load(self.config.load_from))
        self.model.eval()

    def predict(self, batch_embeds, pred_keep_num=100):
        """
        Parameters
        ----------
            batch_embeds: (batch_size=1, token_num, feature_size)
            pred_keep_num: int
        Returns
        -------
            batch_pred: (batch_size, relation_num, object_num, object_num)
            pred_rels: [[sub_id, obj_id, rel_id], ...]
        """
        if not isinstance(batch_embeds, torch.Tensor):
            batch_embeds = torch.asarray(batch_embeds)
        batch_embeds = batch_embeds.to(torch.float32).to(self.device)
        attention_mask = batch_embeds.new_ones((batch_embeds.shape[0], batch_embeds.shape[1]))
        batch_pred = self.model.forward(batch_embeds, attention_mask)
        for idx_i in range(batch_pred.shape[2]):
            batch_pred[:, :, idx_i, idx_i] = -9999
        batch_pred = batch_pred.sigmoid()

        pred_rels = []
        _, topk_indices = torch.topk(
            batch_pred.reshape(
                [
                    -1,
                ]
            ),
            k=pred_keep_num,
        )

        # subject, object, relation
        for index in topk_indices:
            pred_relation = index // (batch_pred.shape[2] ** 2)
            index_subject_object = index % (batch_pred.shape[2] ** 2)
            pred_subject = index_subject_object // batch_pred.shape[2]
            pred_object = index_subject_object % batch_pred.shape[2]
            pred = [pred_subject.item(), pred_object.item(), pred_relation.item()]
            pred_rels.append(pred)
        return batch_pred, pred_rels

    def eval(self):
        sum_recall_20 = 0.0
        sum_recall_50 = 0.0
        sum_recall_100 = 0.0
        prog_bar = ProgressBar(len(self.dataloader))
        for batch_idx, batch_data in enumerate(self.dataloader):
            batch_embeds = batch_data[0]
            batch_target = batch_data[1]
            gt_rels = batch_data[2]
            batch_pred, pred_rels = self.predict(batch_embeds)
            this_recall_20 = get_recall_N(batch_pred, batch_target, object_num=20)
            this_recall_50 = get_recall_N(batch_pred, batch_target, object_num=50)
            this_recall_100 = get_recall_N(batch_pred, batch_target, object_num=100)
            sum_recall_20 += this_recall_20.item()
            sum_recall_50 += this_recall_50.item()
            sum_recall_100 += this_recall_100.item()
            prog_bar.update()
        recall_20 = sum_recall_20 / len(self.dataloader)
        recall_50 = sum_recall_50 / len(self.dataloader)
        recall_100 = sum_recall_100 / len(self.dataloader)
        metric = {
            "recall_20": recall_20,
            "recall_50": recall_50,
            "recall_100": recall_100,
        }
        return metric


if __name__ == "__main__":
    # Config
    config = dict(
        dataset=dict(
            data_path="./data/feat_0420.npz",
            is_train=True,
            num_relation_classes=56,
        ),
        dataloader=dict(
            batch_size=4,
        ),
        model=dict(
            pretrained_model_name_or_path="bert-base-uncased",
            load_pretrained_weights=True,
            num_transformer_layer=2,
            input_feature_size=256,
            output_feature_size=768,
            cls_feature_size=512,
            num_relation_classes=56,
            pred_type="attention",
            loss_type="multi_label_ce",
        ),
        optim=dict(
            lr=1e-4,
            weight_decay=0.05,
            eps=1e-8,
            betas=(0.9, 0.999),
            max_norm=0.01,
            norm_type=2,
            lr_scheduler=dict(
                step=[6, 10],
                gamma=0.1,
            ),
        ),
        num_epoch=12,
        output_dir="./work_dirs",
        load_from=None,
    )

    # Train
    config = Config(config)
    trainer = RamTrainer(config)
    last_model_path = trainer.train()

    # Test/Eval
    config.dataset.is_train = False
    config.load_from = last_model_path
    predictor = RamPredictor(config)
    metric = predictor.eval()
    print(metric)