|
import torch |
|
|
|
|
|
class LoadTextTokens(object): |
|
def __init__(self, tokenizer, max_text_len=40, padding='do_not_pad'): |
|
self.tokenizer = tokenizer |
|
self.max_text_len = max_text_len |
|
self.padding = padding |
|
|
|
def descriptions_to_text_tokens(self, target, begin_token): |
|
target_encoding = self.tokenizer( |
|
target, padding=self.padding, |
|
add_special_tokens=False, |
|
truncation=True, max_length=self.max_text_len) |
|
|
|
need_predict = [1] * len(target_encoding['input_ids']) |
|
payload = target_encoding['input_ids'] |
|
if len(payload) > self.max_text_len - 2: |
|
payload = payload[-(self.max_text_len - 2):] |
|
need_predict = payload[-(self.max_text_len - 2):] |
|
|
|
input_ids = [begin_token] + payload + [self.tokenizer.sep_token_id] |
|
|
|
need_predict = [0] + need_predict + [1] |
|
data = { |
|
'text_tokens': torch.tensor(input_ids), |
|
'text_lengths': len(input_ids), |
|
'need_predict': torch.tensor(need_predict), |
|
} |
|
|
|
return data |
|
|
|
def __call__(self, object_descriptions, box_features, begin_token): |
|
text_tokens = [] |
|
text_lengths = [] |
|
need_predict = [] |
|
for description in object_descriptions: |
|
tokens = self.descriptions_to_text_tokens(description, begin_token) |
|
text_tokens.append(tokens['text_tokens']) |
|
text_lengths.append(tokens['text_lengths']) |
|
need_predict.append(tokens['need_predict']) |
|
|
|
text_tokens = torch.cat(self.collate(text_tokens), dim=0).to(box_features.device) |
|
text_lengths = torch.tensor(text_lengths).to(box_features.device) |
|
need_predict = torch.cat(self.collate(need_predict), dim=0).to(box_features.device) |
|
|
|
assert text_tokens.dim() == 2 and need_predict.dim() == 2 |
|
data = {'text_tokens': text_tokens, |
|
'text_lengths': text_lengths, |
|
'need_predict': need_predict} |
|
|
|
return data |
|
|
|
def collate(self, batch): |
|
if all(isinstance(b, torch.Tensor) for b in batch) and len(batch) > 0: |
|
if not all(b.shape == batch[0].shape for b in batch[1:]): |
|
assert all(len(b.shape) == len(batch[0].shape) for b in batch[1:]) |
|
shape = torch.tensor([b.shape for b in batch]) |
|
max_shape = tuple(shape.max(dim=0)[0].tolist()) |
|
batch2 = [] |
|
for b in batch: |
|
if any(c < m for c, m in zip(b.shape, max_shape)): |
|
b2 = torch.zeros(max_shape, dtype=b.dtype, device=b.device) |
|
if b.dim() == 1: |
|
b2[:b.shape[0]] = b |
|
elif b.dim() == 2: |
|
b2[:b.shape[0], :b.shape[1]] = b |
|
elif b.dim() == 3: |
|
b2[:b.shape[0], :b.shape[1], :b.shape[2]] = b |
|
else: |
|
raise NotImplementedError |
|
b = b2 |
|
batch2.append(b[None, ...]) |
|
else: |
|
batch2 = [] |
|
for b in batch: |
|
batch2.append(b[None, ...]) |
|
return batch2 |
|
else: |
|
raise NotImplementedError |
|
|