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)