import math import copy from contextlib import contextmanager from functools import partial, wraps from pathlib import Path import torch import torch.nn.functional as F from torch import nn, einsum from torch.utils.checkpoint import checkpoint from einops import rearrange, repeat, reduce from einops.layers.torch import Rearrange, Reduce from transformers import BertTokenizer, BertModel from radir import lorentz as L from radir.data.hmclip_mimic_box import HyCoClipMimicBoxDataset_JPG, collate_hmclip_box from radir.data.conditional_dataset_train import SimpleImageTextDataset, simple_collate # NOTE: Import a Cross-Attn Module and PE Module from positional_encodings.torch_encodings import PositionalEncoding3D, PositionalEncoding1D, PositionalEncoding2D from radir.transformer_decoder import TransformerDecoder, TransformerDecoderLayer # from transformer_decoder import TransformerDecoder, TransformerDecoderLayer # NOTE # helper functions def identity(t, *args, **kwargs): return t def exists(val): return val is not None def default(val, d): return val if exists(val) else d @contextmanager def null_context(): yield def max_neg_value(dtype): return -torch.finfo(dtype).max def cast_tuple(t): return t if isinstance(t, (tuple, list)) else (t,) def masked_mean(t, mask, dim = 1, eps = 1e-6): t = t.masked_fill(~mask, 0.) numer = t.sum(dim = dim) denom = mask.sum(dim = dim).clamp(min = eps) return numer / denom def log(t, eps = 1e-20): return torch.log(t + eps) def l2norm(t): return F.normalize(t, dim = -1) def matrix_diag(t): device = t.device i, j = t.shape[-2:] num_diag_el = min(i, j) i_range = torch.arange(i, device = device) j_range = torch.arange(j, device = device) diag_mask = rearrange(i_range, 'i -> i 1') == rearrange(j_range, 'j -> 1 j') diag_el = t.masked_select(diag_mask) return rearrange(diag_el, '(b d) -> b d', d = num_diag_el) # checkpointing helper function def make_checkpointable(fn): @wraps(fn) def inner(*args): input_needs_grad = any([isinstance(el, torch.Tensor) and el.requires_grad for el in args]) if not input_needs_grad: return fn(*args) return checkpoint(fn, *args) return inner # keyword argument helpers def pick_and_pop(keys, d): values = list(map(lambda key: d.pop(key), keys)) return dict(zip(keys, values)) def group_dict_by_key(cond, d): return_val = [dict(),dict()] for key in d.keys(): match = bool(cond(key)) ind = int(not match) return_val[ind][key] = d[key] return (*return_val,) def string_begins_with(prefix, str): return str.startswith(prefix) def group_by_key_prefix(prefix, d): return group_dict_by_key(partial(string_begins_with, prefix), d) def groupby_prefix_and_trim(prefix, d): with_prefix, without_prefix = group_dict_by_key(partial(string_begins_with, prefix), d) trimmed = {k[len(prefix):]: v for k, v in with_prefix.items()} return trimmed, without_prefix def safe_normalize(feat, max_norm=10.0, eps=1e-6): # 支持 (B,D) 或 (B,1,D) orig_shape = feat.shape feat = feat.reshape(-1, feat.size(-1)) norm = feat.norm(dim=-1, keepdim=True) # (B,1) scale = (max_norm / (norm + eps)).clamp(max=1.0) feat = feat * scale return feat.reshape(orig_shape) # main clip class class RADIR(nn.Module): def __init__( self, image_encoder = None, text_encoder = None, tokenizer = None, fusion_module = None, # NOTE: Pass in a Transformer Decoder Module dim_text = 768, dim_image = 768, dim_latent = 768, text_seq_len = 256, text_has_cls_token = False, text_pad_id = 0, text_causal_mask = False, text_eos_id = None, text_encode_without_mask = False, visual_image_size = 224, visual_has_cls_token = False, channels = 3, use_all_token_embeds = False, decoupled_contrastive_learning = False, use_mlm = False, text_ssl_loss_weight = 0.05, use_visual_ssl = False, visual_ssl = None, embed_dim = 512, image_ssl_loss_weight = 0.05, multiview_loss_weight = 0.1, curv_init: float = 1.0, learn_curv: bool = True, entail_weight: float = 0.3, local_aperture_thresh: float = 1.2, global_aperture_thresh: float = 0.7, use_repo_ce_loss: bool = True ): super().__init__() #assert use_all_token_embeds or (visual_has_cls_token or text_has_cls_token), 'CLS token must be included on both vision and text transformers if you are not using fine-grained contrastive learning loss' # store some parameters for access self.dim_text = dim_text self.dim_image = dim_image self.dim_latent = dim_latent self.image_size = visual_image_size # instantiate text transformer self.text_pad_id = text_pad_id self.text_has_cls_token = text_has_cls_token self.text_seq_len = text_seq_len self.text_encode_without_mask = text_encode_without_mask # whether to pass in text mask to text encoder self.text_causal_mask = text_causal_mask self.text_eos_id = text_eos_id assert not (text_causal_mask and not exists(text_eos_id)), 'text EOS token id must be given if using causal mask in text transformer' self.text_transformer = text_encoder # instantiate image transformer self.visual_has_cls_token = visual_has_cls_token self.visual_transformer = image_encoder # NOTE: Pass in a Transformer Decoder Module and a Positional-Encoding Module if exists(fusion_module): self.fusion_module = fusion_module else: decoder_layer = TransformerDecoderLayer(d_model=dim_latent, nhead=8, normalize_before=True) decoder_norm = nn.LayerNorm(dim_latent) self.fusion_module = TransformerDecoder(decoder_layer=decoder_layer, num_layers=4, norm=decoder_norm) self.pos_embedding = PositionalEncoding2D(dim_latent)(torch.zeros(1, 24, 24, dim_latent)) # 1 H W Dim # 修改形状为1 (HW) Dim self.pos_embedding = rearrange(self.pos_embedding, 'b h w d -> b (h w) d') # NOTE # text ssl self.use_mlm = use_mlm self.text_ssl_loss_weight = text_ssl_loss_weight if use_mlm else 0 # image ssl self.use_visual_ssl = use_visual_ssl or exists(visual_ssl) self.image_ssl_loss_weight = image_ssl_loss_weight if use_visual_ssl else 0 self.to_text_latent = nn.Linear(dim_text, dim_latent, bias = False) self.to_visual_latent = nn.Linear(dim_image, dim_latent, bias = False) self.to_fused_latent = nn.Linear(dim_latent, dim_latent, bias = False) # Initialize curvature parameter. Hyperboloid curvature will be `-curv`. self.curv = nn.Parameter( torch.tensor(curv_init).log(), requires_grad=learn_curv ) # When learning the curvature parameter, restrict it in this interval to # prevent training instability. self._curv_minmax = { "max": math.log(curv_init * 10), "min": math.log(curv_init / 10), } # temperature self.temperature = nn.Parameter(torch.tensor(1.)) self.use_repo_ce_loss = use_repo_ce_loss self.logit_scale = nn.Parameter(torch.tensor(1 / 0.07).log()) # 与HycoClip一致的可学习温度 self.entail_weight = entail_weight self.global_aperture_thresh = global_aperture_thresh self.local_aperture_thresh = local_aperture_thresh self.visual_alpha = nn.Parameter(torch.tensor(embed_dim**-0.5).log()) self.textual_alpha = nn.Parameter(torch.tensor(embed_dim**-0.5).log()) # from https://arxiv.org/abs/2111.07783 (FILIP paper) self.use_all_token_embeds = use_all_token_embeds # proposed in https://arxiv.org/abs/2110.06848 (DCL) and https://arxiv.org/abs/2110.11316 (CLOOB) self.decoupled_contrastive_learning = decoupled_contrastive_learning self.multiview_loss_weight = multiview_loss_weight if tokenizer != None: self.tokenizer=tokenizer else: raise ValueError('Tokenzier is not defined.') def load(self, path,allow_partial_load = True): print(f'** MODEL ** Load Checkpoint from {path}') path = Path(path) assert path.exists() pkg = torch.load(path, map_location='cpu', weights_only=False) model_dict = self.state_dict() if 'model' in pkg: pkg_state_dict = pkg['model'] # 假设 pkg['model'] 是模型的状态字典 else: pkg_state_dict = pkg if allow_partial_load: # 检查差异 unexpected_state_dict = [k for k in pkg_state_dict.keys() if k not in model_dict.keys()] missing_state_dict = [k for k in model_dict.keys() if k not in pkg_state_dict.keys()] unmatchd_state_dict = [k for k, v in pkg_state_dict.items() if k in model_dict.keys() and v.shape != model_dict[k].shape] # 加载部分参数 state_dict = {k: v for k, v in pkg_state_dict.items() if k in model_dict.keys() and v.shape == model_dict[k].shape} model_dict.update(state_dict) self.load_state_dict(model_dict) # 可以选择输出具体的参数情况 # print('** MODEL ** The following parameters are UNEXPECTED in checkpoint:\n') # print(unexpected_state_dict) # print('** MODEL ** The following parameters are MISSING in checkpoint:\n') # print(missing_state_dict) # print('** MODEL ** The following parameters have DIFFERENT SHAPES in checkpoint:\n') # print(unmatchd_state_dict) # ('** MODEL ** The following parameters are LOADED in:\n') # print(state_dict.keys()) else: self.load_state_dict(pkg_state_dict) def tokenize(self, prompt): text_tokens=self.tokenizer(prompt, return_tensors="pt", padding="max_length", truncation=True, max_length=512).to('cuda') return text_tokens def token_embedding(self,input_ids): input_shape = input_ids.size() batch_size, seq_length = input_shape if hasattr(self.text_transformer.embeddings, "token_type_ids"): print("hahatrue") buffered_token_type_ids = self.text_transformer.embeddings.token_type_ids[:, :seq_length] buffered_token_type_ids_expanded = buffered_token_type_ids.expand(batch_size, seq_length) token_type_ids = buffered_token_type_ids_expanded text_embeddings = self.text_transformer.embeddings(input_ids = input_ids, token_type_ids = token_type_ids) return text_embeddings def freeze_text_encoder(self): for param in self.text_transformer.parameters(): param.requires_grad = False for param in self.to_text_latent.parameters(): param.requires_grad = False def open_text_encoder(self): for param in self.text_transformer.parameters(): param.requires_grad = True for param in self.to_text_latent.parameters(): param.requires_grad = True def freeze_vision_encoder(self): for param in self.visual_transformer.parameters(): param.requires_grad = False for param in self.to_visual_latent.parameters(): param.requires_grad = False def open_vision_encoder(self): for param in self.visual_transformer.parameters(): param.requires_grad = True for param in self.to_visual_latent.parameters(): param.requires_grad = True def open_fusion_module(self): for param in self.fusion_module.parameters(): param.requires_grad = True for param in self.to_fused_latent.parameters(): param.requires_grad = True def freeze_fusion_module(self): for param in self.fusion_module.parameters(): param.requires_grad = False for param in self.to_fused_latent.parameters(): param.requires_grad = False def get_log_curv(self): return torch.clamp( self.curv, min=self._curv_minmax["min"], max=self._curv_minmax["max"], ) def get_curv(self): with torch.cuda.amp.autocast(enabled=False): curv = torch.exp(self.get_log_curv().float()) return torch.clamp(curv, min=1e-4) def project_img(self, image_feats): image_feats = image_feats * self.visual_alpha.exp() with torch.cuda.amp.autocast(enabled=False): curv = self.get_curv() image_feats = L.exp_map0(image_feats.float(), curv) return image_feats def project_txt(self, text_feats): text_feats = text_feats * self.textual_alpha.exp() with torch.cuda.amp.autocast(enabled=False): curv = self.get_curv() text_feats = L.exp_map0(text_feats.float(), curv) return text_feats def mask_average_pool(self, x, mask): """ 对 token 特征做 masked average pooling Args: x: [B, T, D] (例如 word_feats: [8,128,512]) mask: [B, Tm] 或 [B, K, Tm] 目前的情况通常是 [B,111] Returns: pooled: [B, K, D] """ import torch import torch.nn.functional as F # x: [B, T, D] T = x.size(1) m = mask # [B, Tm] -> [B, 1, Tm] if m.dim() == 2: m = m.unsqueeze(1) # 对齐最后一维到 T（补 0 或截断） Tm = m.size(-1) if Tm < T: m = F.pad(m, (0, T - Tm), value=0) # 右侧补0 elif Tm > T: m = m[..., :T] # einsum 用 float dtype m = m.to(dtype=x.dtype) # 避免全0 mask 导致除0 den = m.sum(dim=-1, keepdim=True).clamp_min(1.0) # [B, K, T] x [B, T, D] -> [B, K, D] return torch.einsum('bkt,btd->bkd', m, x) / den def _unpack_text_batch(self, text): if isinstance(text, dict): input_ids = text.get('caption_ids', text.get('input_ids')) token_type_ids = text.get('token_type_ids', None) attention_mask = text.get('attention_mask', None) else: input_ids = text.input_ids token_type_ids = getattr(text, 'token_type_ids', None) attention_mask = getattr(text, 'attention_mask', None) return input_ids, token_type_ids, attention_mask def forward( self, text, # B image, # B C H W D device, return_loss = True, is_condition = False, modal_embedding = False, phrases_bboxes_mask=None, phrase_text_mask=None, modal_indexs = None, return_latents=True, # 🟢 显式声明 **kwargs ): if is_condition: B, device = text.input_ids.shape[0], device # derive text mask text_mask =text.attention_mask # get encoded text text_args = (text.input_ids,text.attention_mask) if not self.text_encode_without_mask: text_args = (*text_args, text_mask) text_embeddings = self.text_transformer(text.input_ids, attention_mask = text.attention_mask ) enc_text = text_embeddings[0] # B 512(token_len) 768 if image.ndim == 3: # [H, W, ?] 的情况先 unsqueeze image = image.unsqueeze(0) if image.ndim == 4: # B×C×H×W if image.shape[1] != 1: image = image.mean(dim=1, keepdim=True) elif image.ndim == 5: # B×C×D×H×W if image.shape[1] != 1: image = image.mean(dim=1, keepdim=True) else: raise ValueError(...) # NOTE modal_indexs = modal_indexs.view(-1) # (B * local_B) 重复之后，condition就能和uncon的情况一样了 enc_image = self.visual_transformer(image, return_encoded_tokens=True,modal_embedding=modal_embedding,modal_indexs=modal_indexs,is_condition=True) # B*N 24 24 24 512(vis_dim) image_embeds = rearrange(enc_image, 'b t d -> t b d') text_embeds = enc_text[:,0,:] # B 768 text_latents = self.to_text_latent(text_embeds).unsqueeze(0) # 1 B*N 512 image_latents = image_embeds t = image_latents.shape[0] temp_pos_embedding = self.pos_embedding[:,:t, :] # 1 t 512 pos = repeat(temp_pos_embedding, 'one d dim -> d (one bn) dim', bn=image_latents.shape[1]).to(device) # 24*24, B*N, 512 fused_latents, _ = self.fusion_module(text_latents, image_latents, pos=pos) # 1 B*N 512 fused_latents = fused_latents[0] fused_latents = self.to_fused_latent(fused_latents) # B*N 512 norms = torch.norm(fused_latents, dim=-1) # print('fused_latents范数:', '最小值:', norms.min(), '最大值:', norms.max(), '存在零:', (norms == 0).any()) if (norms < 1e-6).any(): print('警告：检测到极小范数') # 可选：添加微小噪声避免零范数问题 fused_latents = fused_latents + torch.randn_like(fused_latents) * 1e-8 if torch.isnan(fused_latents).any() or torch.isinf(fused_latents).any(): print('NaN/Inf in vec **before** l2norm') if fused_latents.isnan().any() or fused_latents.isinf().any(): print('NaN/Inf in vec **after** l2norm') if (fused_latents == 0).any(): print('Zero in vec **after** l2norm') # NOTE if return_latents: return text_latents, image_latents, fused_latents, self.temperature.exp() # ——保持与 else 分支一致 —— # 全局文本/图像向量 text_latents_global = l2norm(self.to_text_latent(enc_text[:, 0, :])) # B x D_latent image_global = l2norm(fused_latents.mean(dim=1)) # B x D_latent # 双曲投影与全局蕴含 #_curv = self.curv.exp() with torch.cuda.amp.autocast(enabled=False): _curv = self.get_curv() text_feats = self.project_txt(text_latents_global) image_feats = self.project_img(image_global) _angle = L.oxy_angle(text_feats, image_feats, _curv) _aperture = L.half_aperture(text_feats, _curv) text_image_entailment_loss = torch.clamp(_angle - self.global_aperture_thresh * _aperture, min=0).mean() # 温度缩放（与仓库版相同） #改curv1 #self.logit_scale.data = torch.clamp(self.logit_scale.data, max=math.log(100.0)) with torch.no_grad(): self.logit_scale.clamp_(max=math.log(100.0)) #改curv2 scale = self.logit_scale.exp() # B×B logits（-双曲距离），对称 CE targets = torch.arange(image_global.size(0), device=image_global.device) image_logits = -L.pairwise_dist(image_feats, text_feats, _curv) # [B, B] text_logits = -L.pairwise_dist(text_feats, image_feats, _curv) # [B, B] # ABLATION: 去掉 image→text CE，只保留 text→image CE i2t_ce_loss = F.cross_entropy(scale * image_logits, targets) t2i_ce_loss = F.cross_entropy(scale * text_logits, targets) ce_loss = t2i_ce_loss # ce_loss = 0.5 * (i2t_ce_loss + t2i_ce_loss) loss = ce_loss + self.entail_weight * text_image_entailment_loss entailment_loss = text_image_entailment_loss return loss, ce_loss, entailment_loss, text_image_entailment_loss else: #text_mask：文本的注意力掩码（布尔矩阵），形状是 B x seq_len。 #text_args：传入文本编码器的集合，包括输入的 token ids 和注意力掩码。 #text_embeddings：文本编码器的输出，包含每个 token 的向量表示。形状是 B x seq_len x dim_text。 #enc_text：文本编码器的最终输出，通常是取 [CLS] token 的向量表示。形状是 B x seq_len x dim_text。 #text_latents：把整句输出文本压成一个固定长度的向量，代表这句话的整体语义。形状是 B x 512。 #text_embeds: 文本的全局向量表示，取自文本编码器的 [CLS] token 输出。形状是 B x dim_text。 #local_B： 单个图像的局部块数量（local_B）。 #modal_indexs：表示每个图像块所属的模态索引，形状是 B x local_B。 #enc_image：图像编码器的输出，包含每个图像块的向量表示。形状是 (B * local_B) x dim_image。 #image_embeds：把所有图像块的向量表示重新排列成 (HWD) x (B * local_B) x dim_image 的形状。 #image_latents：图像块向量，这些局部向量按顺序排成一串，形状是 (HWD) x (B * local_B) x 512。 #box_image_feats_eu：池化后与文本短语对应的图像块特征，形状是 B x K x D_latent。 #box_text_feats_eu：池化后与图像区域对应的文本短语特征，形状是 B x K x D_latent。 input_ids, token_type_ids, attention_mask = self._unpack_text_batch(text) B = input_ids.size(0) with torch.cuda.amp.autocast(enabled=False): _curv = self.get_curv() if image.ndim == 3: # [H, W, ?] 的情况先 unsqueeze image = image.unsqueeze(0) if image.ndim == 4: # B×C×H×W if image.shape[1] != 1: image = image.mean(dim=1, keepdim=True) elif image.ndim == 5: # B×C×D×H×W if image.shape[1] != 1: image = image.mean(dim=1, keepdim=True) else: raise ValueError(...) # ====== 文本编码 ====== text_embeddings = self.text_transformer( input_ids=input_ids, token_type_ids=token_type_ids, attention_mask=attention_mask ) enc_text = text_embeddings[0] # B×T_txt×dim_text #text_latents_global = self.to_text_latent(enc_text[:, 0, :]) text_latents_global = enc_text[:, 0, :] # 不再过 Linear modality_dict = {'CT': 0, 'CXR': 1} modal_indexs = torch.tensor([modality_dict['CXR']], dtype=torch.long, device=device) modal_indexs = modal_indexs.view(-1) # (B * local_B) 重复之后，condition就能和uncon的情况一样了 # ====== 图像编码（双路）====== # Patch 路：box / cross entailment（保留空间 token 网格） enc_image_patch = self.visual_transformer( image, return_encoded_tokens=True, modal_embedding=modal_embedding, modal_indexs=None, is_condition=True, ) # [B, T_img, D] # Global 路：global CE + global entailment（enc_cls_transformer 整图 CLS） enc_image_global = self.visual_transformer( image, return_encoded_tokens=True, modal_embedding=modal_embedding, modal_indexs=None, is_condition=False, ) # [B, D] if self.use_all_token_embeds: text_embeds = enc_text[:, 1:] if self.text_has_cls_token else enc_text else: text_embeds = enc_text[:, :] if enc_text.ndim == 3 else enc_text text_embeds = text_embeds[:, 0, :] text_latents = l2norm(text_embeds) if enc_image_global.dim() == 2: image_latents = l2norm(enc_image_global) elif enc_image_global.dim() == 3: image_latents = l2norm(enc_image_global[:, 0, :]) else: raise ValueError( f"Unexpected global visual feature shape: {enc_image_global.shape}" ) image_global = image_latents word_feats = enc_text patch_feats = enc_image_patch if phrase_text_mask is not None: if phrase_text_mask.dim() == 2: # B×T_txt phrase_text_mask = phrase_text_mask.unsqueeze(1) # B×1×T_txt box_text_feats_eu = self.mask_average_pool(word_feats, phrase_text_mask) # B×K×dim_latent #patch_feats = rearrange(patch_feats_bn, '(b n) t d -> b n t d', n=local_B).mean(dim=1) # B×T_img×dim_latent if phrases_bboxes_mask is not None: if phrases_bboxes_mask.dim() == 2: phrases_bboxes_mask = phrases_bboxes_mask.unsqueeze(1) # B×1×T_img box_image_feats_eu = self.mask_average_pool(patch_feats, phrases_bboxes_mask) # 双曲投影 + 蕴含损失计算（仅当提供 box mask 时） if phrase_text_mask is not None and phrases_bboxes_mask is not None: text_feats = self.project_txt(text_latents_global) # B×D_hyp #image_global = patch_feats.mean(dim=1) # B×D_latent image_feats = self.project_img(image_global) # B×D_hyp box_text_feats = self.project_txt(box_text_feats_eu) # B×K×D_hyp box_image_feats = self.project_img(box_image_feats_eu) # B×K×D_hyp # 后续四项蕴含损失逻辑保持不变 if not return_latents: return text_latents, image_latents, self.temperature.exp() #celoss用，防止被normalize拉爆 image_logits = -L.pairwise_dist(image_feats, text_feats, _curv) # [B, B] text_logits = -L.pairwise_dist(text_feats, image_feats, _curv) # [B, B] #四项蕴含损失 _angle = L.oxy_angle(text_feats, image_feats, _curv) _aperture = L.half_aperture(text_feats, _curv) _box_angle = L.oxy_angle(box_text_feats, box_image_feats, _curv) _box_aperture = L.half_aperture(box_text_feats, _curv) _cross_image_angle = L.oxy_angle(box_image_feats, image_feats.unsqueeze(1), _curv) _box_image_aperture = L.half_aperture(box_image_feats, _curv) _cross_text_angle = L.oxy_angle(box_text_feats, text_feats.unsqueeze(1), _curv) _box_text_aperture = L.half_aperture(box_text_feats, _curv) #全文-全图蕴含 text_image_entailment_loss = torch.clamp(_angle - self.global_aperture_thresh * _aperture, min=0).mean() box_text_image_entailment_loss = torch.clamp(_box_angle - self.global_aperture_thresh * _box_aperture, min=0).mean() cross_image_entailment_loss = torch.clamp(_cross_image_angle - self.local_aperture_thresh * _box_image_aperture, min=0).mean() cross_text_entailment_loss = torch.clamp(_cross_text_angle - self.local_aperture_thresh * _box_text_aperture, min=0).mean() entailment_loss = 0.5 * ( text_image_entailment_loss + box_text_image_entailment_loss + cross_image_entailment_loss + cross_text_entailment_loss ) # 3) 温度缩放 self.logit_scale.data = torch.clamp(self.logit_scale.data, max=4.6052) targets = torch.arange(image_global.size(0), device=image_global.device) scale = self.logit_scale.exp() # ABLATION: 去掉 image→text CE，只保留 text→image CE i2t_ce_loss = F.cross_entropy(scale * image_logits, targets) t2i_ce_loss = F.cross_entropy(scale * text_logits, targets) ce_loss = t2i_ce_loss # ce_loss = 0.5 * (i2t_ce_loss + t2i_ce_loss) loss = ce_loss + self.entail_weight * entailment_loss return loss, ce_loss, entailment_loss, text_image_entailment_loss if __name__ == '__main__': import os import argparse import torch from transformer_maskgit.ctvit import CTViT from transformers import AutoTokenizer, AutoModel, AutoConfig # -------------------------------------------------------- # Args # -------------------------------------------------------- parser = argparse.ArgumentParser(description='Debug RADIR with batch') parser.add_argument('--batch', type=int, default=8) parser.add_argument('--channels', type=int, default=1) parser.add_argument('--height', type=int, default=224) parser.add_argument('--width', type=int, default=224) parser.add_argument('--modal_embedding', action='store_true') # Dataset paths parser.add_argument( '--dataset_root', type=str, default=r"/mnt/data-3/users/zhangzeyu/HMClip/RadIR/dataset" ) parser.add_argument( '--matched_dir', type=str, default=r"/mnt/data-3/users/zhangzeyu/HMClip/RadIR/dataset/processed/matched" ) parser.add_argument( '--caption_file', type=str, default="/mnt/data-3/users/zhangzeyu/HMClip/RadIR/dataset/all.jsonl" # 相对 dataset_root ) # Local HF model path parser.add_argument( '--model_path', type=str, default=r"/mnt/data-3/users/zhangzeyu/HMClip/RadIR/hf_models/microsoft/BiomedVLP-CXR-BERT-specialized" ) args = parser.parse_args() # -------------------------------------------------------- # Device # -------------------------------------------------------- if not torch.cuda.is_available(): raise RuntimeError('需要 GPU 环境（CUDA 可用）。') device = torch.device('cuda') # -------------------------------------------------------- # Debug HF loading # -------------------------------------------------------- from transformers.utils import cached_file try: path = cached_file(args.model_path, "tokenizer_config.json", local_files_only=True) print(">>> HF tokenizer_config.json loaded from:", path) except Exception as e: print(">>> Failed to load tokenizer_config.json:", e) print(">>> Checking config loading...") cfg = AutoConfig.from_pretrained( args.model_path, trust_remote_code=True, local_files_only=True ) print(">>> model_type:", cfg.model_type) print(">>> config path:", cfg._name_or_path) # -------------------------------------------------------- # Load text encoder & tokenizer # -------------------------------------------------------- tokenizer = AutoTokenizer.from_pretrained( args.model_path, trust_remote_code=True, local_files_only=True ) text_model = AutoModel.from_pretrained( args.model_path, trust_remote_code=True, local_files_only=True ).to(device) # -------------------------------------------------------- # Load image encoder (CTViT) # -------------------------------------------------------- image_encoder = CTViT( dim=512, codebook_size=8192, image_size=args.height, patch_size=16, temporal_patch_size=10, spatial_depth=8, temporal_depth=6, cls_depth=4, dim_head=32, heads=8, channels=args.channels ).to(device) # -------------------------------------------------------- # Build RADIR # -------------------------------------------------------- from radir.RadIR import RADIR clip = RADIR( tokenizer=tokenizer, image_encoder=image_encoder, text_encoder=text_model, dim_text=768, dim_image=512, dim_latent=512, use_mlm=False, use_all_token_embeds=False ).to(device) # -------------------------------------------------------- # Build dataset # -------------------------------------------------------- from radir.data.hmclip_mimic_box import ( HyCoClipMimicBoxDataset_JPG, collate_hmclip_box ) from torch.utils.data import DataLoader dataset = HyCoClipMimicBoxDataset_JPG( split="train", tokenizer=tokenizer, dataset_root=args.dataset_root, cig_matched_dir=args.matched_dir, caption_file=args.caption_file, transform=None, # 使用 Dataset 内置 2D + MONAI max_words=128, imsize=224 ) loader = DataLoader( dataset, batch_size=args.batch, shuffle=True, num_workers=0, collate_fn=collate_hmclip_box, ) batch = next(iter(loader)) imgs = batch["imgs"].to(device) # [B, 1, 224, 224] enc_text = { "input_ids": batch["caption_ids"].to(device), "token_type_ids": batch["token_type_ids"].to(device), "attention_mask": batch["attention_mask"].to(device), } phrase_text_mask = batch["phrase_text_mask"].to(device) phrases_bboxes_mask = batch["phrases_bboxes_mask"].to(device) print(f">>> Loaded batch:") print(f" imgs: {imgs.shape}") print(f" input_ids: {enc_text['input_ids'].shape}") print(f" phrase_text_mask: {phrase_text_mask.shape}") print(f" phrases_bboxes_mask: {phrases_bboxes_mask.shape}") # -------------------------------------------------------- # Estimate visual token length T_img # -------------------------------------------------------- with torch.no_grad(): enc_image_tokens = clip.visual_transformer( imgs, return_encoded_tokens=True, modal_embedding=args.modal_embedding, modal_indexs=None, is_condition=False ) T_img = enc_image_tokens.shape[1] print(">>> T_img =", T_img) # -------------------------------------------------------- # Adjust masks # -------------------------------------------------------- phrase_text_mask = phrase_text_mask.unsqueeze(1) # [B,1,T_txt] phrases_bboxes_mask = phrases_bboxes_mask.unsqueeze(1) # [B,1,T_img?] cur_T_img = phrases_bboxes_mask.shape[-1] if cur_T_img != T_img: if cur_T_img > T_img: phrases_bboxes_mask = phrases_bboxes_mask[..., :T_img] else: pad = torch.ones( imgs.size(0), 1, T_img - cur_T_img, device=device ) phrases_bboxes_mask = torch.cat( [phrases_bboxes_mask, pad], dim=-1 ) print(f"[DBG] adjusted phrases_bboxes_mask: {cur_T_img} -> {T_img}") # -------------------------------------------------------- # Modality index (CXR = 1) # -------------------------------------------------------- modality_dict = {'CT': 0, 'CXR': 1} modal_indexs = torch.tensor( [modality_dict['CXR']], dtype=torch.long, device=device ) # -------------------------------------------------------- # Forward # -------------------------------------------------------- print(">>> Forward (return_loss=True)") losses = clip( text=enc_text, image=imgs, device=device, is_condition=False, return_latents=True, return_loss=True, modal_embedding=args.modal_embedding, phrase_text_mask=phrase_text_mask, phrases_bboxes_mask=phrases_bboxes_mask, modal_indexs=modal_indexs ) loss, ce_loss, entailment_loss, text_image_entailment_loss = losses print(f"loss_total: {float(loss):.6f}") print(f"loss_ce(image<->text): {float(ce_loss):.6f}") print(f"loss_entail_global: {float(entailment_loss):.6f}") print(f"loss_entail_image_to_text: {float(text_image_entailment_loss):.6f}")