added in multiple source descriptions

Author: Lisa Dunlap

configs/DomainNet/test_aug.yaml

@@ -0,0 +1,38 @@
+EXP:
+  ADVICE_METHOD: "ClipMLP"
+  WANDB_SILENT: False
+  PROJ: "DomainNetMini_LADS_Replication"
+  SEED: 0
+  TEXT_PROMPTS: [['a realistic photo of a {}.'], ['a painting of a {}.'], ['clipart of a {}.']]
+  NEUTRAL_TEXT_PROMPTS: ['a sketch of a {}', 'a pencil drawing of a {}.', 'a drawing of a {}.']
+  AUGMENTATION: 'DirectionalMulti'
+  EPOCHS: 400
+  LOG_NN: True
+  ENSAMBLE: False
+  
+  
+DATA:
+  DATASET: "DomainNetMini"
+  LOAD_CACHED: True
+  SAVE_PATH: "vit14_clip.pth"
+  BATCH_SIZE: 256
+
+METHOD:
+  MODEL:
+    NUM_LAYERS: 1
+    DOM_WEIGHT: 1.0
+    LR: 0.0001
+    CHECKPOINT: 'checkpoint/mlp_simple.pth'
+    CHECKPOINT_NAME: 'DomainNetMini-mlp-directional'
+    RESUME: False
+  USE_DOM_GT: True
+  APPLY_TRANSFORMATION: False
+
+AUGMENTATION:
+  MODEL:
+    LR: 0.0001
+    WEIGHT_DECAY: 0.005
+    NUM_LAYERS: 1
+  EPOCHS: 50
+  GENERIC: False
+  ALPHA: 0.5

methods/augmentations.py

@@ -600,6 +600,219 @@ def load_checkpoint(self, net, path):
         print(f"...loaded checkpoint with acc {checkpoint['acc']}")
         return net
 
+from clip_utils import get_domain_text_embs
+
+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 DirectionalMulti(Augment):
+
+    def __init__(self, cfg, image_features, labels, group_labels, domain_labels, filenames, text_features, val_image_features, val_labels, val_group_labels,val_domain_labels, val_filenames):
+        super().__init__(cfg, image_features, labels, group_labels, domain_labels, filenames, text_features)
+        source_embeddings, target_embeddings = get_domain_text_embs(self.model, cfg, self.neutral_prompts, self.prompts, self.class_names)
+        # target_embeddings is size (num_domains, num_classes, emb_size)
+        # source_embeddings is size (num_source_domain_descriptions, num_classes, emb_size)
+        source_embeddings /= source_embeddings.norm(dim=-1, keepdim=True)
+        target_embeddings /= target_embeddings.norm(dim=-1, keepdim=True)
+        self.source_embeddings = source_embeddings.cuda().float()
+        self.target_embeddings = target_embeddings.cuda().float()
+        dataset = EmbeddingDataset(self.cfg, self.image_features, self.labels, self.group_labels, self.domain_labels)
+        self.dataset = dataset
+        self.train_loader = torch.utils.data.DataLoader(dataset, batch_size=cfg.DATA.BATCH_SIZE, shuffle=True)
+
+        val_dataset = EmbeddingDataset(self.cfg, val_image_features, val_labels, val_group_labels, val_domain_labels)
+        self.val_dataset = val_dataset
+        self.val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=cfg.DATA.BATCH_SIZE, shuffle=True)
+
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.nets = []
+        self.net_checkpoints = []
+        self.uid = uuid.uuid4()
+        if self.cfg.DATA.DATASET == 'ColoredMNISTBinary':
+            text_embs = zeroshot_classifier([[f'a photo of the number "{c}"'] for c in self.class_names], self.model, model_type=self.cfg.EXP.IMAGE_FEATURES)
+        else:
+            text_embs = zeroshot_classifier([[f"a photo of a {c}"] for c in self.class_names], self.model, model_type=self.cfg.EXP.IMAGE_FEATURES)
+        
+        self.class_text_embs = text_embs.float().cuda()
+        print("text emb shape ", self.class_text_embs.shape)
+
+        for i in range(len(self.prompts)):
+            print(f"Training network for {self.prompts[i]}")
+            self.train_network(i)
+
+    def directional_loss_builder(self):
+        """
+        CLIP directional loss from gan NADA paper. Ensures that the difference in
+        image embeddings is similar to the difference in text embeddings of the 
+        source and target domain.
+        """
+        def custom_loss(predictions, labels, targets):
+            total_sum = None
+            delta_i = predictions - labels
+            ctr = 0
+            for i, delta_tt in zip(delta_i, targets): 
+                ctr += 1
+                if total_sum == None: 
+                    numerator = torch.dot(i, delta_tt)
+                    denominator = torch.norm(i) * torch.norm(delta_tt)
+                    total_sum = 1 - (numerator/denominator)
+                else: 
+                    total_sum += 1 - (torch.dot(i, delta_tt)/ (torch.norm(i) * torch.norm(delta_tt)))
+            return total_sum / ctr
+        return custom_loss
+
+    @staticmethod
+    def get_class_logits(outputs, class_embs):
+        outputs_norm = outputs / outputs.norm(dim=-1, keepdim=True) 
+        return torch.matmul(outputs_norm, class_embs)
+
+    def train_network(self, num_net): 
+        net = MLP(hidden_dim=self.cfg.AUGMENTATION.MODEL.HIDDEN_DIM, input_dim=self.dataset.embedding_dim)
+        self.nets.append(net.cuda())
+        self.net_checkpoints.append("")
+
+        self.optimizer = AdamW(self.nets[num_net].parameters(), lr=self.cfg.AUGMENTATION.MODEL.LR, weight_decay=self.cfg.AUGMENTATION.MODEL.WEIGHT_DECAY)
+        self.directional_loss = DirectionLoss(self.cfg.AUGMENTATION.LOSS_TYPE)
+        # self.directional_loss = self.directional_loss_builder()
+        self.class_consistency_loss = nn.CrossEntropyLoss(weight=self.dataset.class_weights.cuda())
+
+        if self.cfg.AUGMENTATION.CLIP_NN_LOSS:
+            self.clip_nn_loss = nn.CrossEntropyLoss()
+
+        self.nets[num_net].train()
+        
+        best_train_loss, best_epoch = 10000, 0
+        for epoch in range(self.cfg.AUGMENTATION.EPOCHS):
+            train_metrics = self.training_loop(self.train_loader, num_net, epoch, phase='train')
+            val_metrics = self.training_loop(self.val_loader, num_net, epoch, phase='val')
+            # val_metrics = self.eval_loop(num_net, epoch)
+            if val_metrics['val loss'] < best_train_loss:
+                    best_train_loss = val_metrics['val loss']
+                    best_epoch = epoch
+                    self.net_checkpoints[num_net] = self.save_checkpoint(best_train_loss, epoch, num_net)
+
+        wandb.summary[f"{self.prompts[num_net]} best epoch"] = best_epoch
+        wandb.summary[f"{self.prompts[num_net]} best train_loss"] = best_train_loss
+        print(f"==> loading checkpoint {self.net_checkpoints[num_net]} at epoch {best_epoch} with loss {best_train_loss}")
+        self.nets[num_net] = self.load_checkpoint(self.nets[num_net], self.net_checkpoints[num_net])
+
+    def get_direction_vectors(self, img_embs, labels, num_net):
+        """
+        Returns the direction vectors for the image embeddings by taking the source
+        embedding that is most similar to each image embedding and subtracting if from the target.
+        """
+        dir_vectors = []
+        for (im, l) in zip(img_embs, labels):
+            prod = im @ self.source_embeddings[:,l,:].T
+            _, source_idx = torch.max(prod, dim=0)
+            diff = self.target_embeddings[num_net][l] - self.source_embeddings[source_idx][l]
+            if diff.norm() == 0:
+                print(diff)
+            dir_vectors.append(diff)
+        diffs = torch.stack(dir_vectors)
+        diffs /= diffs.norm(dim=-1, keepdim=True)
+        return diffs
+
+    def training_loop(self, loader, num_net, epoch, phase='train'):
+        if phase == 'train':
+            self.nets[num_net].train()
+        else:
+            self.nets[num_net].eval()
+        train_directional_loss, train_class_loss, train_loss, total = 0, 0, 0, 0
+        with torch.set_grad_enabled(phase == 'train'):
+            for i, (inp, cls_target, cls_group, dom_target) in enumerate(loader):
+                inp, cls_target= inp.cuda().float(), cls_target.cuda().long()
+                cls_outputs = self.nets[num_net](inp)
+                text_diffs = self.get_direction_vectors(inp, cls_target, num_net)
+                im_diffs = cls_outputs - inp
+                # print(text_diffs.shape, im_diffs.shape)
+                # print(torch.min(text_diffs, dim=0), torch.max(text_diffs, dim=0))
+                # text_diffs -= text_diffs.min(dim=-1, keepdim=True)
+                # im_diffs -= im_diffs.min(dim=-1, keepdim=True)
+                # text_diffs /= text_diffs.norm(dim=-1, keepdim=True)
+                # im_diffs /= im_diffs.norm(dim=-1, keepdim=True)
+                # compute directional loss
+                # directional_loss = self.directional_loss(inp, cls_outputs/cls_outputs.norm(dim=-1, keepdim=True), text_diffs)
+                directional_loss = self.directional_loss(im_diffs / im_diffs.norm(dim=-1, keepdim=True), text_diffs).mean()
+                # print(directional_loss)
+                cls_logits = self.get_class_logits(cls_outputs, self.class_text_embs)
+                cls_consist = self.class_consistency_loss(cls_logits, cls_target)
+                loss = self.alpha * directional_loss + (1 - self.alpha) * cls_consist
+                train_class_loss += (1 - self.alpha) * cls_consist.item()
+                train_directional_loss += self.alpha * directional_loss.item()
+
+                if phase == 'train':
+                    self.optimizer.zero_grad()
+                    loss.backward(retain_graph=True)
+                    self.optimizer.step()
+                
+                train_loss += loss.item()
+
+                total += cls_target.size(0)
+                progress_bar(i, len(loader), 'Loss: %.3f'% (train_loss/(i+1)))
+
+        metrics = {f"{phase} class loss": train_class_loss/(i+1), f"{phase} directional loss": train_directional_loss/(i+1), f"{phase} loss": train_loss/(i+1), "epoch": epoch}
+        wandb.log(metrics)
+        return metrics
+
+    def augment_single(self, img_embedding, label): 
+        keep = img_embedding
+        if self.cfg.AUGMENTATION.INCLUDE_ORIG_TRAINING:
+            output = [keep]
+        else:
+            output = []
+        img_embedding = torch.tensor(img_embedding)
+        img_embedding = img_embedding.type(torch.float32)
+        img_embedding = img_embedding.cuda()
+        img_embedding /= img_embedding.norm(dim=-1, keepdim=True) 
+        for net in self.nets: 
+            o = net(img_embedding)
+            o /= o.norm(dim=-1, keepdim=True) 
+            
+            o = o.detach().cpu().numpy()
+            output.append(o)
+            wandb.log({"cos sim:": distance.cosine(o, img_embedding.cpu())})
+        # output = self.net(img_embedding)
+        # val = torch.tensor(output)
+        return list(np.array(output))
+
+    def save_checkpoint(self, acc, epoch, num_net):
+        checkpoint_dir = os.path.join("./checkpoint", self.cfg.DATA.DATASET)
+        if not os.path.exists(checkpoint_dir):
+            os.makedirs(checkpoint_dir)
+        path = f'./checkpoint/{self.cfg.DATA.DATASET}/{self.prompts[num_net]}-{self.cfg.EXP.SEED}-{self.uid}.pth'
+        print(f'Saving checkpoint with acc {acc} to {path}...')
+        state = {
+            "acc": acc,
+            "epoch": epoch,
+            "net": self.nets[num_net].state_dict()
+        }
+        torch.save(state, path)
+        # wandb.save(path)
+        return path
+
+    def load_checkpoint(self, net, path):
+        checkpoint = torch.load(path)
+        net.load_state_dict(checkpoint['net'])
+        print(f"...loaded checkpoint with acc {checkpoint['acc']}")
+        return net
+
 class BiasDirectional(Directional):
     """
     This implements the similar directional loss as the directional class, but routes examples