File size: 10,679 Bytes
998bb30 |
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 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 |
import torch
import torchvision.transforms as transforms
import torchvision.transforms.functional as TF
import numpy as np
from typing import Union, List
import os
from imagebind import data
from imagebind.models import imagebind_model
from imagebind.models.imagebind_model import ModalityType
class DirectionLoss(torch.nn.Module):
def __init__(self, loss_type='mse'):
super(DirectionLoss, self).__init__()
self.loss_type = loss_type
self.loss_func = {
'mse': torch.nn.MSELoss,
'cosine': torch.nn.CosineSimilarity,
'mae': torch.nn.L1Loss
}[loss_type]()
def forward(self, x, y):
if self.loss_type == "cosine":
return 1. - self.loss_func(x, y)
return self.loss_func(x, y)
class ImageBindLoss(torch.nn.Module):
def __init__(self, device, lambda_direction=1., lambda_naive=0., direction_loss_type='cosine', pretrained=True, class_names=None):
super(ImageBindLoss, self).__init__()
self.device = device
# Load ImageBind model
self.model = imagebind_model.imagebind_huge(pretrained=pretrained)
self.model.eval() # Eval mode, but gradients will still flow for adversarial attacks
self.model.to(device)
# Loss components (unchanged from CLIPLoss)
self.direction_loss = DirectionLoss(direction_loss_type)
self.patch_direction_loss = torch.nn.CosineSimilarity(dim=2)
self.cosine_sim = torch.nn.CosineSimilarity(dim=1)
self.lambda_naive = lambda_naive
self.lambda_direction = lambda_direction
self.src_text_features = None
self.target_text_features = None
self.angle_loss = torch.nn.L1Loss()
self.text_class_features = dict()
self.image_class_features = dict()
self.predicted_classes = class_names
self.mse = torch.nn.MSELoss()
def tokenize(self, strings: list):
"""Tokenize text strings using ImageBind's data.load_and_transform_text."""
return data.load_and_transform_text(strings, self.device)
def encode_text(self, tokens: torch.Tensor) -> torch.Tensor:
"""Encode tokenized text using ImageBind model, without gradients."""
inputs = {ModalityType.TEXT: tokens}
with torch.no_grad(): # Text features are typically fixed, no gradients needed
embeddings = self.model(inputs)
return embeddings[ModalityType.TEXT]
def preprocess(self, x: torch.Tensor) -> torch.Tensor:
"""Preprocess tensor images from [-1, 1] to ImageBind's expected format."""
x = (x + 1) / 2 # Convert from [-1, 1] to [0, 1]
x = TF.resize(x, 224, interpolation=TF.InterpolationMode.BICUBIC)
x = TF.center_crop(x, 224)
x = TF.normalize(
x,
mean=(0.48145466, 0.4578275, 0.40821073),
std=(0.26862954, 0.26130258, 0.27577711)
)
return x
def encode_images(self, images: torch.Tensor) -> torch.Tensor:
"""Encode images using ImageBind model, preserving gradients."""
images = self.preprocess(images).to(self.device)
inputs = {ModalityType.VISION: images}
embeddings = self.model(inputs) # No torch.no_grad() to allow gradients
return embeddings[ModalityType.VISION]
def get_text_features(self, class_str: Union[str, List[str]]) -> torch.Tensor:
"""Retrieve precomputed text features for given classes."""
if isinstance(class_str, str):
class_str = [class_str]
text_features = [self.text_class_features[i] for i in class_str]
return torch.stack(text_features, dim=0)
def precompute_text_features(self, class_str: Union[str, List[str]], templates=None, norm: bool = True) -> None:
"""Precompute and store text features for given classes."""
if isinstance(class_str, str):
class_str = [class_str]
for classes in class_str:
if classes not in self.text_class_features:
template_text = self.compose_text_with_templates(classes, templates)
tokens = self.tokenize(template_text)
text_features = self.encode_text(tokens)
text_features = text_features.mean(dim=0)
if norm: # ImageBind embeddings are already normalized, but kept for consistency
text_features /= text_features.norm(dim=-1, keepdim=True)
self.text_class_features[classes] = text_features
def precompute_image_features(self, images: dict, norm: bool = True) -> None:
"""Precompute and store image features for given classes."""
for classes in images:
if classes not in self.image_class_features:
class_images = images[classes]
all_images = class_images
batch_size = 16
num_batches = (all_images.size(0) + batch_size - 1) // batch_size
batched_image_features = []
for i in range(num_batches):
start_idx = i * batch_size
end_idx = min((i + 1) * batch_size, all_images.size(0))
batch = all_images[start_idx:end_idx]
batch_features = self.encode_images(batch)
batched_image_features.append(batch_features.detach()) # Detach to save memory
image_features = torch.cat(batched_image_features, dim=0)
image_features = image_features.mean(dim=0)
if norm:
image_features /= image_features.norm(dim=-1, keepdim=True)
self.image_class_features[classes] = image_features
def get_image_features(self, img: torch.Tensor, norm: bool = True) -> torch.Tensor:
"""Get image features with optional normalization."""
image_features = self.encode_images(img)
if norm:
image_features /= image_features.clone().norm(dim=-1, keepdim=True)
return image_features
def compute_text_direction(self, source_class: Union[str, List[str]], target_class: Union[str, List[str]], broadcast=False) -> torch.Tensor:
"""Compute the direction between source and target text features."""
source_features = self.get_text_features(source_class)
target_features = self.get_text_features(target_class)
if broadcast:
text_direction = (target_features.T.unsqueeze(0) - source_features.unsqueeze(-1))
else:
text_direction = (target_features - source_features)
text_direction /= text_direction.norm(dim=-1, keepdim=True)
return text_direction
def compose_text_with_templates(self, text: str, templates=None) -> list:
"""Compose text with templates (assumes templates are provided)."""
if templates is None:
templates = ["{}"] # Default template if none provided
return [template.format(text) for template in templates]
def clip_directional_loss(self, src_img: torch.Tensor, source_class: Union[str, List[str]], target_img: torch.Tensor, target_class: Union[str, List[str]], negative_class: Union[str, List[str]]) -> torch.Tensor:
"""Compute directional loss using ImageBind embeddings."""
self.target_direction = self.compute_text_direction(source_class, target_class).detach()
self.negative_direction = self.compute_text_direction(source_class, negative_class, broadcast=True).detach()
src_encoding = self.get_image_features(src_img)
target_encoding = self.get_image_features(target_img)
edit_direction = (target_encoding - src_encoding)
edit_direction /= (edit_direction.clone().norm(dim=-1, keepdim=True) + 1e-7)
logit_target = self.cosine_sim(self.target_direction, edit_direction)
logit_negative = self.cosine_sim(self.negative_direction, edit_direction.unsqueeze(-1))
pp = torch.exp(logit_target)
pn = torch.sum(torch.exp(logit_negative), dim=-1)
p = pp / (pp + pn)
return -torch.log(p).mean()
def clip_class_loss(self, target_img: torch.Tensor, target_class: Union[str, List[str]], negative_class: Union[str, List[str]]) -> torch.Tensor:
"""Compute class loss using ImageBind embeddings."""
text_features = self.get_text_features(target_class)
negative_text_features = self.get_text_features(negative_class)
image_features = self.get_image_features(target_img)
logit_target = self.cosine_sim(text_features, image_features)
logit_negative = self.cosine_sim(
negative_text_features.unsqueeze(0).expand(len(target_img), -1, -1).permute(0, 2, 1),
image_features.unsqueeze(-1)
)
pp = torch.exp(logit_target)
pn = torch.sum(torch.exp(logit_negative), dim=-1)
p = pp / (pp + pn)
return -torch.log(p).mean()
def forward(self, src_img: torch.Tensor, source_class: Union[str, List[str]], target_img: torch.Tensor, target_class: Union[str, List[str]], negative_class: Union[str, List[str]], texture_image: torch.Tensor = None) -> torch.Tensor:
"""Forward pass combining naive and directional losses."""
clip_loss = 0.0
# if self.lambda_naive:
# naive_loss = self.clip_class_loss(target_img, target_class, negative_class)
# clip_loss += self.lambda_naive * naive_loss
if self.lambda_direction:
direction_loss = self.clip_directional_loss(src_img, source_class, target_img, target_class, negative_class)
clip_loss += self.lambda_direction * direction_loss
return clip_loss
def predict(self, img: torch.Tensor) -> torch.Tensor:
"""Predict class probabilities using ImageBind embeddings."""
image_features = self.get_image_features(img)
text_features = self.get_text_features(self.predicted_classes)
similarity = (100.0 * image_features @ text_features.T).softmax(dim=-1)
return similarity
if __name__ == '__main__':
device = 'cuda' if torch.cuda.is_available() else 'cpu'
imagebind_loss = ImageBindLoss(device)
# Test with small batch
src_img = torch.zeros(2, 3, 32, 32, device=device)
tgt_img = torch.zeros(2, 3, 32, 32, device=device)
loss = imagebind_loss(src_img, ['dog', 'cat'], tgt_img, ['boy', 'cat'], ['bird', 'fish'])
print(f"Loss with batch size 2: {loss.item()}")
# Test with larger batch
src_img = torch.zeros(64, 3, 32, 32, device=device)
tgt_img = torch.zeros(64, 3, 32, 32, device=device)
loss = imagebind_loss(src_img, 'dog', tgt_img, 'cat', 'bird')
print(f"Loss with batch size 64: {loss.item()}") |