<jupyter_start><jupyter_text>**NT-Xent Begins**<jupyter_code># data_dir = '/data/majie/Datasets/TinyDataset/hymenoptera_data'
# ant = osp.join(data_dir, 'train', 'ants', 'VietnameseAntMimicSpider.jpg')
# bee = osp.join(data_dir, 'train', 'bees', '39747887_42df2855ee.jpg')
# def load_image(file_path):
#     image = Image.open(file_path).convert('RGB')
# #     image = np.asarray(image)
#     return image # PIL
# transform_crop = transforms.Compose([transforms.RandomCrop(224), transforms.ToTensor()])
# transform_color = transforms.Compose([transforms.Resize((224, 224)), 
#                                       transforms.ColorJitter(1, 0.1, 0.05, 0.0), 
#                                       transforms.ToTensor()])
# ant_crop = transform_crop(load_image(ant))
# ant_color = transform_color(load_image(ant)) # torch.Size([3, 224, 224])
# bee_crop = transform_crop(load_image(bee))
# bee_color = transform_color(load_image(bee)) # torch.Size([3, 224, 224])
# ant1 = ant_crop.view(1, -1)
# ant2 = ant_color.view(1, -1)
# bee1 = bee_crop.view(1, -1)
# bee2 = bee_color.view(1, -1)
animals = [ant1, ant2, bee1, bee2]
# 构造 L2 distance 矩阵， 计算 每张 图片 与 所有图片的距离 N * N
# s1 = [F.pairwise_distance(animals[0], animals[x]) for x in range(4)]
# s2 = [F.pairwise_distance(animals[1], animals[x]) for x in range(4)]
# s3 = [F.pairwise_distance(animals[2], animals[x]) for x in range(4)]
# s4 = [F.pairwise_distance(animals[3], animals[x]) for x in range(4)]
S = torch.Tensor([[torch.cosine_similarity(animals[i], animals[x]) for x in range(4)] for i in range(4)])
# S = torch.Tensor([[F.pairwise_distance(animals[i], animals[x]) for x in range(4) if x != i] for i in range(4)])
S
# 让同一张图的不同view在latent space里靠近，不同图的view在latent space里远离，通过NT-Xent实现
Loss12 = S[0][1] / sum(S[0])
Loss21 = S[1][0] / sum(S[1])
Loss34 = S[2][3] / sum(S[2])
Loss43 = S[3][2] / sum(S[3])
Loss12 + Loss21
Loss34 + Loss43
def L_distance(k):
    # k: [1,.., N]
    return S[2*k - 1-1][2*k-1] / torch.sum(S[2 *k-1-1]) + S[2*k - 1][2*k - 1-1] / torch.sum(S[2*k -1])
L_distance(1)
L_distance(2)
loss = 0
for k in range(1, 3):
    loss += L_distance(k)
loss / 4<jupyter_output><empty_output><jupyter_text>**Apply NT-Xent in multi domains**<jupyter_code>def load_image(file_path):
    image = Image.open(file_path).convert('RGB')
#     image = np.asarray(image)
    return image # PIL
data_dir = '/data/majie/Datasets/Domain-adaption/PACS/data'
doms = ['cartoon', 'photo', 'sketch']
doms_test = 'art_painting'
doms_test = 0
dom_datasets = {}
data_sizes = {}
image_datasets = {}
image_datasets_test = {}
device = torch.device("cuda:3" if torch.cuda.is_available() else "cpu")
device
for i, dom in enumerate(doms):
    image_datasets[dom] = {x: datasets.ImageFolder(os.path.join(data_dir, dom, x)) for x in ['train', 'val', 'test']}
test_dom_dataset = datasets.ImageFolder(osp.join(data_dir, 'art_painting', 'test'), transforms.Compose([
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ]))
test_dom_dataset
test_dom_dataloader = DataLoader(test_dom_dataset, batch_size=16, shuffle=False)
test_data = test_dom_dataloader.__iter__().__next__()
inputs = test_data[0].to(device)
out = model_ft(inputs, 0.99)
test_data[1]
model_ft = resnet18()
model_ft = model_ft.to(device)
model_ft.load_state_dict(torch.load('./models/best_model_adversarial_contrastive.tar')['state'])
model_ft.eval()
# acc = 0
# for images, labels in test_dom_dataloader:
#     images = images.to(device)
#     labels = labels.to(device)
    
#     outputs = model_ft(images, 0.8)[1]
#     _, preds_classes = torch.max(outputs, 1)
#     print('preds_classes', preds_classes)
#     acc += torch.sum(preds_classes == labels.data)
# print('test on art domain, accuracy is: ', acc)
train, val = [], []
for idx, dom in enumerate(doms):
    train_label = [image_datasets[dom]['train'].imgs[i] + (idx,) for i in range(len(image_datasets[dom]['train'].imgs))]
    train.append(train_label)
    val_label = [image_datasets[dom]['val'].imgs[i] + (idx,) for i in range(len(image_datasets[dom]['val'].imgs))]
    val.append(val_label)
train = np.vstack(np.array(train))
val = np.vstack(np.array(val))
train[1839]
# image_datasets
class GenericDataset(Dataset):
    def __init__(self, dataset_root='', stage='train', random_sized_crop=False, num_imgs_per_cat=None):
        super(GenericDataset, self).__init__()
        
        # if split is not None, return the spcific split data
        # else return all data in source domains
        
        self.image_file_names = []
        self.stage = stage
        
        self.dataset = {'train': [], 'val': []}
        self.data = []
        
        if self.stage == 'test': # test domain
            
            dom = doms[doms_test] 
            # test split is the whole data in its domain
            self.data.extend(image_datasets[dom]['test'].imgs)
            
        else: # training domain
           
            self.dataset['train'] = train
            self.dataset['val'] = val
            
            self.data = self.dataset[self.stage]
       
            
          
    
    def __getitem__(self, index):
        # 1. return img - label
        image, class_label, dom_label = self.data[index]
        
        return image, class_label, dom_label

    
    def __len__(self):
        
        return len(self.data)
dataset = GenericDataset()
dataloader = DataLoader(dataset,batch_size=4)
data = dataloader.__iter__().__next__()
len(data)
data_transform = {
    'train': {
        'crop': transforms.Compose([
            transforms.RandomResizedCrop(224),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ]),
        'color': transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.RandomHorizontalFlip(),
            transforms.ColorJitter(1, 0.1, 0.05, 0.0), 
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ])
    },
    'val': {
        'crop': transforms.Compose([
            transforms.CenterCrop((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ]),
        'color': transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ColorJitter(1, 0.1, 0.05, 0.0), 
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ])
    },
    'test': transforms.Compose([
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])
}
class DataLoader_custom(object):
    def __init__(self,
                 dataset,
                 stage='train',
                 batch_size=8,
                 unsupervised=True,
                 epoch_size=None,
                 num_workers=8,
                 shuffle=True):
        
        self.dataset = dataset
        self.shuffle = shuffle
        self.epoch_size = epoch_size if epoch_size is not None else len(dataset)
        self.batch_size = batch_size
        self.unsupervised = unsupervised
        self.num_workers = num_workers
        
        self.stage = stage

        mean_pix  = [0.485, 0.456, 0.406]
        std_pix = [0.229, 0.224, 0.225]
        
        
        self.transform = data_transform
     

    def get_iterator(self, epoch=0):
        rand_seed = epoch * self.epoch_size
        random.seed(rand_seed)
        # if in unsupervised mode define a loader function that given the
        # index of an image it returns the 4 rotated copies of the image
        # plus the label of the rotation, i.e., 0 for 0 degrees rotation,
        # 1 for 90 degrees, 2 for 180 degrees, and 3 for 270 degrees.
        
        if self.stage == 'test':
            def _load_function(idx):
#                 idx = idx % len(self.dataset)
                img0, class_label, dom_label = self.dataset[idx]
                img0 = load_image(img0)
                image = self.transform[self.stage](img0)
               
                class_label = torch.LongTensor(int(class_label))
                dom_label = torch.LongTensor(int(dom_label))
                
#                 print('--------------------------')

                return image, (class_label, dom_label)

            def _collate_fun(batch):
                return default_collate(batch)
        
        else:
            def _load_function(idx):
                idx = idx % len(self.dataset)
                img0, class_label, dom_label = self.dataset[idx]

                img0 = load_image(img0)
                rotated_imgs = [
                    self.transform[self.stage]['crop'](img0),
                    self.transform[self.stage]['color'](img0),
#                     transform(rotate_img(img0,  90)),
#                     transform(rotate_img(img0, 180)),
#                     transform(rotate_img(img0, 270))
                ]
                labels = [(int(class_label), int(dom_label)) for _ in range(2)]
                rotation_labels = torch.LongTensor(labels)
                return torch.stack(rotated_imgs, dim=0), rotation_labels

            def _collate_fun(batch):
                batch = default_collate(batch)
                assert(len(batch)==2)
#                 print('batch[1]', batch[1])
                batch_size, rotations, channels, height, width = batch[0].size()
                batch[0] = batch[0].view([batch_size*rotations, channels, height, width])
                batch[1] = batch[1].view((batch_size*rotations, 2))
                return batch


        tnt_dataset = tnt.dataset.ListDataset(elem_list=range(self.epoch_size),
            load=_load_function)
        data_loader = tnt_dataset.parallel(batch_size=self.batch_size,
            collate_fn=_collate_fun, num_workers=self.num_workers,
            shuffle=self.shuffle)
        return data_loader

    def __call__(self, epoch=0):
        return self.get_iterator(epoch)

    def __len__(self):
        return self.epoch_size / self.batch_size
dataloader = {x: DataLoader_custom(GenericDataset(stage=x), stage=x) for x in ['train', 'val', 'test']}
data_size = {x: len(GenericDataset(stage=x)) for x in ['train', 'val', 'test']}
data_size
len(dataloader['train']())
class DataPrefetcher():
    def __init__(self, loader):
        self.loader = iter(loader)
        self.stream = torch.cuda.Stream()
        # With Amp, it isn't necessary to manually convert data to half.
        # if args.fp16:
        #     self.mean = self.mean.half()
        #     self.std = self.std.half()
        self.preload()

    def preload(self):
        try:
            self.batch = next(self.loader)
        except StopIteration:
            self.batch = None
            return
        with torch.cuda.stream(self.stream):
            for k in range(len(self.batch)):
                if k != 'meta':
                    self.batch[k] = self.batch[k].to(device=device, non_blocking=True)

            # With Amp, it isn't necessary to manually convert data to half.
            # if args.fp16:
            #     self.next_input = self.next_input.half()
            # else:
            #     self.next_input = self.next_input.float()

    def next(self):
        torch.cuda.current_stream().wait_stream(self.stream)
        batch = self.batch
        self.preload()
        return batch
prefetcher = DataPrefetcher(dataloader['train']())
image, label = prefetcher.next()

image.size()<jupyter_output><empty_output><jupyter_text>**TODO:** check the dataloader order as above according to Paper: A Simple Framework for Contrastive Learning of Visual Representations**REFER:** [adverserial training](https://zhuanlan.zhihu.com/p/50710267)<jupyter_code>def contrastive_distance(k):
    # k: [1,.., N]
    return -torch.log(torch.exp(S[2*k - 1-1][2*k-1]) / torch.sum(torch.exp(S[2 *k-1-1]))\
+ torch.exp(S[2*k - 1][2*k - 1-1]) / torch.sum(torch.exp(S[2*k -1])))<jupyter_output><empty_output><jupyter_text>**TODO  alpha setting:** [refer](https://github.com/ShichengChen/Domain-Adversarial-Training-of-Neural-Networks/blob/master/mnist.ipynb)<jupyter_code>def train_model(model, criterion, optimizer, scheduler, num_epochs=50):
    best_model_wts = copy.deepcopy(model.state_dict())
    dataloader = {x: DataLoader_custom(GenericDataset(stage=x), stage=x) for x in ['train', 'val', 'test']}
    data_size = {x: len(GenericDataset(stage=x)) for x in ['train', 'val', 'test']}
    best_acc = 0.0
    alpha = 0.0


    for epoch in range(num_epochs):
        print('Epoch {}/{}'.format(epoch, num_epochs))
        print('-' * 20)
        
        i = 0
     
        for phase in ['train', 'val']:
            len_dataloader = len(dataloader[phase](epoch))
            dom_data_size = data_size[phase]
            if phase == 'train':
                total_steps = num_epochs * len_dataloader
                start_steps = epoch * len_dataloader
                p = float(i + start_steps) / total_steps
                alpha = 2. / (1. + np.exp(-10 * p)) - 1
                model.train()
            else:
                model.eval()

            running_loss = 0.0
            running_corrects = 0

            ## train each dom, sum over all losses, and backward the total loss

            doms_data_size = 0
            # prefetcher = DataPrefetcher(dataloader['train']())
            dom_data_loader = DataPrefetcher(dataloader[phase](epoch))
            
            
            
            batch = dom_data_loader.next()
            while batch is not None:
                inputs, labels  = batch

#             for inputs, labels in dom_data_loader.next():
#                 inputs = inputs.to(device)
#                 labels = labels.to(device)

                label_class = labels[:, 0]
                label_dom = labels[:, 1]

                # zero the parameter gradients
                optimizer.zero_grad()

                with torch.set_grad_enabled(phase == 'train'):
                    outputs = model(inputs, alpha) # 0: contrastive_feature, 1: class_label, 2: dom_label

                    batch_data_size = inputs.size(0)
                    
                    k_len = batch_data_size / 2 + 1
                    
                    contrastive_features = outputs[0].view(batch_data_size, 1, 128)
                    S = torch.Tensor([[torch.cosine_similarity(contrastive_features[i], contrastive_features[x])\
                                       for x in range(batch_data_size)] for i in range(batch_data_size)])

                    contrastive_loss = 0
                    for k in range(1, int(k_len)):
                        contrastive_distance =  -torch.log(torch.exp(S[2*k - 1-1][2*k-1]) / torch.sum(torch.exp(S[2 *k-1-1]))\
+ torch.exp(S[2*k - 1][2*k - 1-1]) / torch.sum(torch.exp(S[2*k -1])))
                        contrastive_loss += contrastive_distance

                    contrastive_loss = contrastive_loss.to(device)
                    loss_class = criterion(outputs[1], label_class)
                    loss_dom = criterion(outputs[2], label_dom)

                    total_loss = 0.6 * (contrastive_loss + loss_class) + 0.4 * loss_dom

                    _, preds_classes = torch.max(outputs[1], 1)
                    _, preds_dom = torch.max(outputs[2], 1)


                    if phase == 'train':
                        i += 1
                        total_loss.backward()
                        optimizer.step()
                    batch = dom_data_loader.next()

                running_loss += total_loss.item() * inputs.size(0)
                running_corrects += torch.sum(preds_classes == label_class.data)
            
            if phase == 'train':
                scheduler.step()

            epoch_loss = running_loss / dom_data_size / 2
            epoch_acc = running_corrects.double() / dom_data_size / 2

            print('{} Loss: {:.4f} Acc: {:.4f}'.format(
                phase, epoch_loss, epoch_acc))

            # deep copy the model
            if phase == 'val' and epoch_acc > best_acc:
                best_acc = epoch_acc
                best_model_wts = copy.deepcopy(model.state_dict())
            


    print('Best val Acc: {:4f}'.format(best_acc))

    # load best model weights
    model.load_state_dict(best_model_wts)

    if not os.path.exists('./models'):
        os.makedirs('./models')

    outfile = os.path.join('./models', 'best_model_adversarial_contrastive_0.4.tar')
    torch.save({'state': best_model_wts, 'alpha': alpha}, outfile)
    return model
model = resnet18()
# model = nn.DataParallel(model, device_ids=[0,1,2,3])

model = model.to(device)
criterion = nn.CrossEntropyLoss()
# Observe that all parameters are being optimized
optimizer = optim.Adam(model.parameters(), lr=1e-3)

# Decay LR by a factor of 0.1 every 7 epochs
exp_lr_scheduler = lr_scheduler.StepLR(optimizer, step_size=7, gamma=0.1)
device
model = train_model(model, criterion, optimizer, exp_lr_scheduler, num_epochs=20)
test_dom_dataset = datasets.ImageFolder(osp.join(data_dir, 'art_painting', 'test'), transforms.Compose([
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ]))
test_dom_dataloader = DataLoader(test_dom_dataset, batch_size=16, shuffle=False)
acc = 0
for images, labels in test_dom_dataloader:
    images = images.to(device)
    labels = labels.to(device)
    
    outputs = model(images, 0.8)[1]
    _, preds_classes = torch.max(outputs, 1)
    
    acc += torch.sum(preds_classes == labels.data)
    print('Accuarcy ', acc)
print('test on art domain, accuracy is: ', acc)
# multi GPU train, to load saved weights
def load_network(network):
    save_path = os.path.join('./model',name,'net_%s.pth'%opt.which_epoch)
    state_dict = torch.load(save_path)
    # create new OrderedDict that does not contain `module.`
    from collections import OrderedDict
    new_state_dict = OrderedDict()
    for k, v in state_dict.items():
        namekey = k[7:] # remove `module.`
        new_state_dict[namekey] = v
    # load params
    network.load_state_dict(new_state_dict)
    return network<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.13392760082617675, 'fraction_numerical': 0.030220676160452223, 'mean_word_length': 2.3229185479501537, 'pattern_counts': {'":': 0, '<': 17, '<?xml version=': 0, '>': 18, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '4', 'dir': 'jupyter-structured-clean-dedup', 'id': '345883', 'license': 'MIT', 'n_tokens_mistral': 6434, 'n_tokens_neox': 5996, 'n_words': 1207, 'path': 'main.ipynb', 'repo_name': 'Shadow-JM/SimCLR', 'size': '18398'}