File size: 7,559 Bytes
320c5b2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 | import torch
import torch.nn as nn
import torch.nn.functional as F
from model.clip import build_model
from utils.simple_tokenizer import SimpleTokenizer as _Tokenizer
from .layers import FPN, Projector, TransformerDecoder, MultiTaskProjector
_tokenizer = _Tokenizer()
def tokenize_text(texts, context_length, truncate=False):
if isinstance(texts, str):
texts = [texts]
sot_token = _tokenizer.encoder["<|startoftext|>"]
eot_token = _tokenizer.encoder["<|endoftext|>"]
all_tokens = [[sot_token] + _tokenizer.encode(text) + [eot_token] for text in texts]
result = torch.zeros(len(all_tokens), context_length, dtype=torch.long)
for i, tokens in enumerate(all_tokens):
if len(tokens) > context_length:
if truncate:
tokens = tokens[:context_length]
tokens[-1] = eot_token
else:
raise RuntimeError(f"Input {texts[i]} is too long for context length {context_length}")
result[i, :len(tokens)] = torch.tensor(tokens)
return result
class CROG(nn.Module):
def __init__(self, cfg):
super().__init__()
# Flags for ablation study
self.use_contrastive = cfg.use_contrastive
self.use_pretrained_clip = cfg.use_pretrained_clip
self.use_grasp_masks = cfg.use_grasp_masks
self.word_len = cfg.word_len
# Vision & Text Encoder
clip_model = torch.jit.load(cfg.clip_pretrain,
map_location="cpu").eval()
print(f"Load pretrained CLIP: {self.use_pretrained_clip}")
self.backbone = build_model(clip_model.state_dict(), cfg.word_len, self.use_pretrained_clip).float()
# Multi-Modal FPN
self.neck = FPN(in_channels=cfg.fpn_in, out_channels=cfg.fpn_out)
# Decoder
if self.use_contrastive:
print("Use contrastive learning module")
# Decoder
self.decoder = TransformerDecoder(num_layers=cfg.num_layers,
d_model=cfg.vis_dim,
nhead=cfg.num_head,
dim_ffn=cfg.dim_ffn,
dropout=cfg.dropout,
return_intermediate=cfg.intermediate)
else:
print("Disable contrastive learning module")
if self.use_grasp_masks:
# Projector
print("Use grasp masks")
self.proj = MultiTaskProjector(cfg.word_dim, cfg.vis_dim // 2, 3)
else:
print("Disable grasp masks")
self.proj = Projector(cfg.word_dim, cfg.vis_dim // 2, 3)
def _prepare_word(self, word, device):
if isinstance(word, str):
word = [word]
if isinstance(word, (list, tuple)):
word = tokenize_text(list(word), self.word_len, True)
if not torch.is_tensor(word):
raise TypeError(f"CROG expects token ids or raw text, got {type(word)}")
return word.to(device=device, dtype=torch.long, non_blocking=True)
def forward(self, img, word, mask=None, grasp_qua_mask=None, grasp_sin_mask=None, grasp_cos_mask=None, grasp_wid_mask=None):
'''
img: b, 3, h, w
word: b, words
word_mask: b, words
mask: b, 1, h, w
'''
word = self._prepare_word(word, img.device)
# padding mask used in decoder
pad_mask = torch.zeros_like(word).masked_fill_(word == 0, 1).bool()
# vis: C3 / C4 / C5
# word: b, length, 1024
# state: b, 1024
vis = self.backbone.encode_image(img)
word, state = self.backbone.encode_text(word)
# print(f"[DEBUG] word shape: {word.shape}, state shape: {state.shape}")
# b, 512, 26, 26 (C4)
fq = self.neck(vis, state)
b, c, h, w = fq.size()
if self.use_contrastive:
fq = self.decoder(fq, word, pad_mask)
fq = fq.reshape(b, c, h, w)
if self.use_grasp_masks:
# b, 1, 104, 104
pred, grasp_qua_pred, grasp_sin_pred, grasp_cos_pred, grasp_wid_pred = self.proj(fq, state)
if self.training:
# resize mask
if pred.shape[-2:] != mask.shape[-2:]:
mask = F.interpolate(mask, pred.shape[-2:], mode='nearest').detach()
grasp_qua_mask = F.interpolate(grasp_qua_mask, grasp_qua_pred.shape[-2:], mode='nearest').detach()
grasp_sin_mask = F.interpolate(grasp_sin_mask, grasp_sin_pred.shape[-2:], mode='nearest').detach()
grasp_cos_mask = F.interpolate(grasp_cos_mask, grasp_cos_pred.shape[-2:], mode='nearest').detach()
grasp_wid_mask = F.interpolate(grasp_wid_mask, grasp_wid_pred.shape[-2:], mode='nearest').detach()
# Ratio Augmentation
total_area = mask.shape[2] * mask.shape[3]
coef = 1 - (mask.sum(dim=(2,3)) / total_area)
# Generate weight
weight = mask * 0.5 + 1
loss = F.binary_cross_entropy_with_logits(pred, mask, weight=weight)
grasp_qua_loss = F.smooth_l1_loss(grasp_qua_pred, grasp_qua_mask)
grasp_sin_loss = F.smooth_l1_loss(grasp_sin_pred, grasp_sin_mask)
grasp_cos_loss = F.smooth_l1_loss(grasp_cos_pred, grasp_cos_mask)
grasp_wid_loss = F.smooth_l1_loss(grasp_wid_pred, grasp_wid_mask)
# @TODO adjust coef of different loss items
total_loss = loss + grasp_qua_loss + grasp_sin_loss + grasp_cos_loss + grasp_wid_loss
loss_dict = {}
loss_dict["m_ins"] = loss.item()
loss_dict["m_qua"] = grasp_qua_loss.item()
loss_dict["m_sin"] = grasp_sin_loss.item()
loss_dict["m_cos"] = grasp_cos_loss.item()
loss_dict["m_wid"] = grasp_wid_loss.item()
# loss = F.binary_cross_entropy_with_logits(pred, mask, reduction="none").sum(dim=(2,3))
# loss = torch.dot(coef.squeeze(), loss.squeeze()) / (mask.shape[0] * mask.shape[2] * mask.shape[3])
return (pred.detach(), grasp_qua_pred.detach(), grasp_sin_pred.detach(), grasp_cos_pred.detach(), grasp_wid_pred.detach()), (mask, grasp_qua_mask, grasp_sin_mask, grasp_cos_mask, grasp_wid_mask), total_loss, loss_dict
else:
return (pred.detach(), grasp_qua_pred.detach(), grasp_sin_pred.detach(), grasp_cos_pred.detach(), grasp_wid_pred.detach()), (mask, grasp_qua_mask, grasp_sin_mask, grasp_cos_mask, grasp_wid_mask)
else:
# b, 1, 104, 104
pred = self.proj(fq, state)
if self.training:
# resize mask
if pred.shape[-2:] != mask.shape[-2:]:
mask = F.interpolate(mask, pred.shape[-2:],
mode='nearest').detach()
loss = F.binary_cross_entropy_with_logits(pred, mask)
loss_dict = {}
loss_dict["m_ins"] = loss.item()
loss_dict["m_qua"] = 0
loss_dict["m_sin"] = 0
loss_dict["m_cos"] = 0
loss_dict["m_wid"] = 0
return (pred.detach(), None, None, None, None), (mask, None, None, None, None), loss, loss_dict
else:
return pred.detach(), mask
|