{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Introduction\nThis is an **UNOFFICIAL** implementation of *Asymmetric Tri-Training for Unsupervised Domain Adaptation (ATDA)* from the following paper \n>K. Saito, Y. Ushiku, and T. Harada, “Asymmetric Tri-training for Unsupervised Domain Adaptation,” in Proceedings of the 34th International Conference on Machine Learning, PMLR, Jul. 2017, pp. 2988–2997. Accessed: Jul. 05, 2023. [Online]. Available: https://proceedings.mlr.press/v70/saito17a.html\n\nThe authors of the paper released their official implementation at https://github.com/ksaito-ut/atda/tree/master. However, only the code of two of the experiments in the paper are publicized: domain adaptation from SVHN to MNIST, and the reverse. Besides, the code was written in TensorFlow 1.3.0, which was very old. Thus, based on the source code from the aforementioned repo, I rewrite the ATDA algorithm to the PyTorch version and implement the experiment of bidirectional domain adaptation for MNIST to USPS (which are not implemented in their paper).\n\n## Disclaimer\nThe detailed model architecture for the experiment is unavailable. Although the authors provided the model architecture in the supplementary of the paper, they did not clarify where they place the dropout layers, for example. They only mentioned `We also used dropout in our experiment.` in the caption of Figure 2 of the supplementary. Also, the padding mode, values in the 2D convolutional layers are not provided. Thus, this implementation may not be accurate.\n\n## Results\n| Task | USPS -> MNIST | MNIST -> USPS |\n|----------|---------:|---------:|\n| Source-only | 85.2 | 94.4 |\n| Adapted | 93.1 | 97.2 |","metadata":{}},{"cell_type":"markdown","source":"# Prepare Libraries","metadata":{}},{"cell_type":"markdown","source":"## Install Necessary Libraries\nUncomment the cell if you did not install these libraries before.  \n1. `umap-learn`: To visualize embeddings from the source and target domains by UMAP","metadata":{}},{"cell_type":"code","source":"!pip install umap-learn","metadata":{"execution":{"iopub.status.busy":"2023-08-25T07:53:10.047279Z","iopub.execute_input":"2023-08-25T07:53:10.047654Z","iopub.status.idle":"2023-08-25T07:53:21.613829Z","shell.execute_reply.started":"2023-08-25T07:53:10.047622Z","shell.execute_reply":"2023-08-25T07:53:21.612565Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Import Libraries","metadata":{}},{"cell_type":"code","source":"import os\nimport h5py\nimport pickle\nimport random\nfrom itertools import cycle\n\nfrom tqdm.notebook import tqdm\nimport numpy as np\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\nimport torch.backends.cudnn as cudnn\nfrom torchvision import datasets, transforms\nfrom torch.utils.data import Dataset, Subset, DataLoader, random_split\nfrom matplotlib import pyplot as plt\nimport umap","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2023-08-25T07:53:21.616596Z","iopub.execute_input":"2023-08-25T07:53:21.617286Z","iopub.status.idle":"2023-08-25T07:53:21.624587Z","shell.execute_reply.started":"2023-08-25T07:53:21.617243Z","shell.execute_reply":"2023-08-25T07:53:21.623561Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Global variables and preparation","metadata":{}},{"cell_type":"code","source":"def create_exp_directory(exp_dir):\n    if not os.path.exists(exp_dir):\n        os.makedirs(exp_dir)\n        os.makedirs(f'{exp_dir}/umap')\n        print(f'Directory {exp_dir} created.')\n        print(f'Directory {exp_dir}/umap created.')\n    \n    if not os.path.exists(f'{exp_dir}/umap'): \n        os.makedirs(f'{exp_dir}/umap')  \n        print(f'Directory {exp_dir}/umap created.')   \n\n        \nEXP_DIR = './exps'\nDEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nSOURCE_DIR = './data/MNIST'\nTARGET_DIR = '/kaggle/input/usps-dataset/usps.h5'\nS_TRANSFORM = transforms.Compose([transforms.ToTensor(),\n                                  transforms.Normalize(mean=[0.5, ], std=[0.5, ]),\n                                 ],)\nT_TRANSFORM = transforms.Compose([transforms.ToTensor(),\n                                  transforms.Pad([6, 6], fill=0.),\n                                  transforms.Normalize(mean=[0.5, ], std=[0.5, ]),\n                                 ],)\n\ncreate_exp_directory(EXP_DIR)\nprint(f'Currently using device: {DEVICE}')","metadata":{"execution":{"iopub.status.busy":"2023-08-25T07:53:21.626467Z","iopub.execute_input":"2023-08-25T07:53:21.626840Z","iopub.status.idle":"2023-08-25T07:53:21.640860Z","shell.execute_reply.started":"2023-08-25T07:53:21.626804Z","shell.execute_reply":"2023-08-25T07:53:21.639814Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Training Hyperparameters","metadata":{}},{"cell_type":"code","source":"cudnn.benchmark = True # Speed up training\n\nN_CLASSES = 10\nLR = 0.01\nWC_LAMBDA = 0.01 # Lambda for weight constraint\nN_ITERS = 30\nN_STEPS = 1500\nBATCH_SIZE = 128\nTHRESHOLD = 0.95","metadata":{"execution":{"iopub.status.busy":"2023-08-25T07:53:21.645155Z","iopub.execute_input":"2023-08-25T07:53:21.645444Z","iopub.status.idle":"2023-08-25T07:53:21.655602Z","shell.execute_reply.started":"2023-08-25T07:53:21.645419Z","shell.execute_reply":"2023-08-25T07:53:21.654737Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Datasets Preparation","metadata":{}},{"cell_type":"code","source":"class USPS(Dataset):\n    def __init__(self, data_path, train=True, transform=None, target_transform=None, download=True):\n        self.train = train\n        self.transform = transform\n        self.target_transform = target_transform\n        \n        with h5py.File(data_path, 'r') as hf:\n            if train:\n                train = hf.get('train')\n                self.data = train.get('data')[:]\n                self.labels = train.get('target')[:]\n            else:\n                test = hf.get('test')\n                self.data = test.get('data')[:]\n                self.labels = test.get('target')[:]\n        \n    def __len__(self):\n        return len(self.data)\n    \n    def __getitem__(self, index):\n        img = self.data[index].reshape(16, 16, 1)\n        label = self.labels[index].astype(np.compat.long)\n        \n        if self.transform:\n            img = self.transform(img)\n            \n        if self.target_transform:\n            label = self.target_transform(label)\n        \n        return img, label\n\ns_train_dataset = datasets.MNIST(SOURCE_DIR, train=True, transform=S_TRANSFORM, download=True)\nt_train_dataset = USPS(TARGET_DIR, train=True, transform=T_TRANSFORM, download=True)","metadata":{"execution":{"iopub.status.busy":"2023-08-25T07:53:21.657811Z","iopub.execute_input":"2023-08-25T07:53:21.658136Z","iopub.status.idle":"2023-08-25T07:53:21.785831Z","shell.execute_reply.started":"2023-08-25T07:53:21.658104Z","shell.execute_reply":"2023-08-25T07:53:21.784755Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Visualize Data from Both Source and Target Domains","metadata":{}},{"cell_type":"code","source":"random_indices = random.choices(np.arange(len(t_train_dataset)), k=4)\n\n# Visualize the selected images\nplt.figure(figsize=(10, 5))\n\nfor i, index in enumerate(random_indices):\n    s_img = s_train_dataset[index][0].permute(1, 2, 0)\n    t_img = t_train_dataset[index][0].permute(1, 2, 0)\n    \n    plt.subplot(2, 4, i + 1)\n    plt.imshow(s_img, cmap='gray')\n    plt.title(f\"MNIST, label={s_train_dataset[index][1]}\")\n    plt.axis('off')\n    \n    plt.subplot(2, 4, i + 5)\n    plt.imshow(t_img, cmap='gray')\n    plt.title(f\"USPS, label={t_train_dataset[index][1]}\")\n    plt.axis('off')\n\nplt.tight_layout()\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2023-08-25T07:53:21.787117Z","iopub.execute_input":"2023-08-25T07:53:21.787479Z","iopub.status.idle":"2023-08-25T07:53:22.532542Z","shell.execute_reply.started":"2023-08-25T07:53:21.787444Z","shell.execute_reply":"2023-08-25T07:53:22.531588Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_dataset_stats(dataloader):\n    num_channels = 0\n    mean = torch.zeros(3)  # Assuming you are using RGB images (3 channels)\n    std = torch.zeros(3)\n    max_values = torch.zeros(3)\n    min_values = torch.zeros(3)    \n    num_samples = 0\n\n    for images, _ in tqdm(dataloader):\n        batch_size = images.size(0)\n        num_channels = images.size(1)\n\n        # Calculate mean and std for each channel separately\n        if num_channels == 1:\n            images_view = images.view(batch_size, -1)\n            mean += images_view.mean(1).sum(0)\n            std += images_view.std(1).sum(0)\n            max_values = images.max()\n            min_values = images.min()\n        else:\n            images_view = images.view(batch_size, num_channels, -1)\n            mean += images_view.mean(2).sum(0)\n            std += images_view.std(2).sum(0)\n            max_values = images.max()\n            min_values = images.min()\n\n        num_samples += batch_size\n\n    mean /= num_samples\n    std /= num_samples \n    return mean, std, max_values, min_values\n\ns_loader = DataLoader(s_train_dataset, batch_size=BATCH_SIZE)\ns_mean, s_std, s_max, s_min = get_dataset_stats(s_loader)\nprint(f'Source Dataloader: Mean={s_mean}, STD={s_std}, Max={s_max}, Min={s_min}')\n\nt_loader = DataLoader(t_train_dataset, batch_size=BATCH_SIZE)\nt_mean, t_std, t_max, t_min = get_dataset_stats(t_loader)\nprint(f'Target Dataloader: Mean={t_mean}, STD={t_std}, Max={t_max}, Min={t_min}')","metadata":{"execution":{"iopub.status.busy":"2023-08-25T07:53:22.534076Z","iopub.execute_input":"2023-08-25T07:53:22.535152Z","iopub.status.idle":"2023-08-25T07:53:35.983423Z","shell.execute_reply.started":"2023-08-25T07:53:22.535114Z","shell.execute_reply":"2023-08-25T07:53:35.982340Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Prepare Training","metadata":{}},{"cell_type":"markdown","source":"## Model Architrcture\nATDA model proposed by\n>K. Saito, Y. Ushiku, and T. Harada, “Asymmetric Tri-training for Unsupervised Domain Adaptation,” in Proceedings of the 34th International Conference on Machine Learning, PMLR, Jul. 2017, pp. 2988–2997. Accessed: Jul. 05, 2023. [Online]. Available: https://proceedings.mlr.press/v70/saito17a.html\n\n\n![Model Architecture](https://i.imgur.com/vld7ccY.png)","metadata":{}},{"cell_type":"code","source":"class ATDA(nn.Module):\n    def __init__(self):\n        super().__init__()\n        self.f = FeatureExtractor()\n        self.f1 = LabelNet(drop_rate=0.5)\n        self.f2 = LabelNet(drop_rate=0.5)\n        self.ft = TargetNet(drop_rate=0.2)\n\n    def forward(self, x):\n        embeddings = self.f(x)\n        logits_1 = self.f1(embeddings)\n        logits_2 = self.f2(embeddings)\n        logits_t = self.ft(embeddings)\n        return logits_1, logits_2, logits_t\n    \n    def get_embeddings(self, x):\n        return self.f(x)\n\nclass FeatureExtractor(nn.Module):\n    def __init__(self):\n        super().__init__()\n        self.conv1 = nn.Conv2d(1, 32, kernel_size=5)\n        self.conv2 = nn.Conv2d(32, 48, kernel_size=5, bias=False)\n        self.bn_conv2 = nn.BatchNorm2d(48)\n        self.maxpool = nn.MaxPool2d(kernel_size=2)\n        self.relu = nn.ReLU()\n\n    def forward(self, x):\n        x = self.conv1(x)\n        x = self.relu(x)\n        x = self.maxpool(x)\n\n        x = self.conv2(x)\n        x = self.relu(x)\n        x = self.bn_conv2(x)\n        x = self.maxpool(x)\n        return x.flatten(1)\n\nclass LabelNet(nn.Module):\n    def __init__(self, drop_rate):\n        super().__init__()\n        self.fc1 = nn.Linear(768, 100, bias=False)\n        self.bn1 = nn.BatchNorm1d(100)\n        self.fc2 = nn.Linear(100, 100, bias=False)\n        self.bn2 = nn.BatchNorm1d(100)\n        self.fc3 = nn.Linear(100, 10, bias=False)\n\n        self.relu = nn.ReLU()\n        self.dropout = nn.Dropout(p=drop_rate)\n\n    def forward(self, x):\n        x = self.fc1(x)\n        x = self.relu(x)\n        x = self.bn1(x)\n        x = self.dropout(x)\n\n        x = self.fc2(x)\n        x = self.bn2(x)\n        x = self.relu(x)\n        x = self.dropout(x)\n\n        x = self.fc3(x)\n        return x\n    \nclass TargetNet(nn.Module):\n    def __init__(self, drop_rate):\n        super().__init__()\n        self.fc1 = nn.Linear(768, 100, bias=True)\n        self.fc2 = nn.Linear(100, 100, bias=True)\n        self.fc3 = nn.Linear(100, 10, bias=True)\n\n        self.relu = nn.ReLU()\n        self.dropout = nn.Dropout(p=drop_rate)\n\n    def forward(self, x):\n        x = self.fc1(x)\n        x = self.relu(x)\n        x = self.dropout(x)\n\n        x = self.fc2(x)\n        x = self.relu(x)\n        x = self.dropout(x)\n\n        x = self.fc3(x)\n        return x\n\nmodel = ATDA().to(DEVICE)","metadata":{"execution":{"iopub.status.busy":"2023-08-25T07:53:35.985159Z","iopub.execute_input":"2023-08-25T07:53:35.985848Z","iopub.status.idle":"2023-08-25T07:53:36.013828Z","shell.execute_reply.started":"2023-08-25T07:53:35.985811Z","shell.execute_reply":"2023-08-25T07:53:36.012862Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Optimizer and Loss Function\n### Loss function\nBeside the standard *Cross Entropy loss*, the paper introduced a loss function, which I implemented as the `WeightConstraint` abstract below, to force the two heads labeling net to learn differently by the equation\n$$L_{WC}=\\lambda |W_1^TW_2|$$\nwhere $\\lambda$ is a hyperparameter to scale this loss.\n\n> Because we aim to label the target samples with high accuracy, we expect $F_1$ and $F_2$ to classify the samples based on different viewpoints. Therefore, we make a constraint for the weights of $F_1$ and $F_2$ to make their inputs different from each other. We add the term $|W_1^TW_2|$ to the cost function, where $W_1$ and $W_2$ denote fully connected layerweights of $F_1$ and $F_2$, which are first applied to the feature $F(x_i)$.","metadata":{}},{"cell_type":"code","source":"class WeightConstraint(nn.Module):\n    def __init__(self, lambda_):\n        \"\"\"\n        To force two heads f1 and f2 in the model learn different\n        weights so they can provide better pseudo-labels\n        \n        (From the paper): 'Because we aim to label the target samples \n        with high accuracy, we expect F1 and F2 to classify the samples \n        based on different viewpoints. Therefore, we make a constraint \n        for the weights of F1 and F2 to make their inputs different from \n        each other. We add the term jW1T W2j to the cost function, where \n        W1 and W2 denote fully connected layerweights of F1 and F2, which \n        are first applied to the feature F (xi).'\n\n        Args:\n            model1 (nn.Module)\n            model2 (nn.Module)\n            lambda_ (float): \n                lambda value for the weight constraint. \n                Set based on validation split. \n                Paper: 0.01.\n        \"\"\"\n        super(WeightConstraint, self).__init__()\n        self.lambda_ = lambda_\n\n    def forward(self, model1, model2):\n        # Get the learnable parameters of model1 (F_1)\n        w1 = [param for param in model1.parameters() if param.requires_grad]\n        w1 = torch.cat([p.view(-1) for p in w1])\n\n        # Get the learnable parameters of model2 (F_2)\n        w2 = [param for param in model2.parameters() if param.requires_grad]\n        w2 = torch.cat([p.view(-1) for p in w2])\n\n        # Calculate the dot product and take the absolute value\n        dot_product = torch.dot(w1, w2)\n        loss = self.lambda_ * torch.abs(dot_product)\n        return loss","metadata":{"execution":{"iopub.status.busy":"2023-08-25T07:53:36.015219Z","iopub.execute_input":"2023-08-25T07:53:36.015769Z","iopub.status.idle":"2023-08-25T07:53:36.027043Z","shell.execute_reply.started":"2023-08-25T07:53:36.015729Z","shell.execute_reply":"2023-08-25T07:53:36.026044Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"optims = {'fn': optim.SGD(list(model.f.parameters()) + list(model.f1.parameters()) + list(model.f2.parameters()), lr=LR, momentum=0.9),\n          'ft': optim.SGD(list(model.f.parameters()) + list(model.ft.parameters()), lr=LR, momentum=0.9)}\ncriteria = {'ce': nn.CrossEntropyLoss(), 'wc': WeightConstraint(lambda_=WC_LAMBDA)}","metadata":{"execution":{"iopub.status.busy":"2023-08-25T07:53:36.030812Z","iopub.execute_input":"2023-08-25T07:53:36.031157Z","iopub.status.idle":"2023-08-25T07:53:36.043355Z","shell.execute_reply.started":"2023-08-25T07:53:36.031123Z","shell.execute_reply":"2023-08-25T07:53:36.042406Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Training and Evaluation","metadata":{}},{"cell_type":"markdown","source":"## Helper functions for training","metadata":{}},{"cell_type":"code","source":"class CustomImageDataset(Dataset):\n    def __init__(self, X: torch.tensor, y: torch.tensor):\n        self.X = X\n        self.y = y\n    \n    def __len__(self):\n        return self.X.shape[0]\n    \n    def __getitem__(self, index):\n        img = self.X[index]\n        label = self.y[index]\n        return img, label\n    \ndef extract_datasets(dataset):\n    x = []\n    y = []\n    for data, label in dataset:\n        x.append(data)\n        y.append(label)\n    x = torch.stack(x, dim=0)\n    y = torch.tensor(y, dtype=torch.long)\n    return x, y\n\ndef run_pretrain(model, optims, criteria, X, y):\n    model.train()\n    \n    optims['fn'].zero_grad()\n    logits_1, logits_2, _ = model(X)\n    ce_loss_1 = criteria['ce'](logits_1, y)\n    ce_loss_2 = criteria['ce'](logits_2, y)\n    wc_loss = criteria['wc'](model.f1, model.f2)\n    loss_fn = ce_loss_1 + ce_loss_2 + wc_loss\n    loss_fn.backward()\n    optims['fn'].step()\n\n    optims['ft'].zero_grad()\n    _, _, logits_t = model(X)\n    ce_loss_t = criteria['ce'](logits_t, y)\n    ce_loss_t.backward()\n    optims['ft'].step()\n\n    total_loss = (ce_loss_1 + ce_loss_2 + ce_loss_t + wc_loss).item()\n\n    _, predicted_labels_1 = logits_1.max(1)\n    _, predicted_labels_2 = logits_2.max(1)\n    _, predicted_labels_t = logits_t.max(1)\n    acc_f1 = (predicted_labels_1 == y).to(torch.float).mean().item()\n    acc_f2 = (predicted_labels_2 == y).to(torch.float).mean().item()\n    acc_ft = (predicted_labels_t == y).to(torch.float).mean().item()\n    return total_loss, wc_loss, ce_loss_1, ce_loss_2, ce_loss_t, acc_f1, acc_f2, acc_ft\n\ndef run_retrain(model, optims, criteria, X, y, mode):\n    model.train()\n    \n    if mode == 'mixed':\n        optims['fn'].zero_grad()\n        logits_1, logits_2, _ = model(X)\n        ce_loss_1 = criteria['ce'](logits_1, y)\n        ce_loss_2 = criteria['ce'](logits_2, y)\n        wc_loss = criteria['wc'](model.f1, model.f2)\n        loss_fn = ce_loss_1 + ce_loss_2 + wc_loss\n        loss_fn.backward()\n        optims['fn'].step()\n\n        _, predicted_labels_1 = logits_1.max(1)\n        _, predicted_labels_2 = logits_2.max(1)\n        acc_f1 = (predicted_labels_1 == y).to(torch.float).mean().item()\n        acc_f2 = (predicted_labels_2 == y).to(torch.float).mean().item()\n        total_loss = (ce_loss_1 + ce_loss_2 + wc_loss).item()\n        return total_loss, wc_loss, ce_loss_1, ce_loss_2, acc_f1, acc_f2\n    elif mode == 'target':\n        optims['ft'].zero_grad()\n        _, _, logits_t = model(X)\n        ce_loss_t = criteria['ce'](logits_t, y)\n        ce_loss_t.backward()\n        optims['ft'].step() \n\n        _, predicted_labels_t = logits_t.max(1)\n        acc_ft = (predicted_labels_t == y).to(torch.float).mean().item()\n        return acc_ft\n    else:\n        raise ValueError(f'Unknown mode: {mode}. Options: [\"mixed\", \"target\"]')\n\ndef run_eval(model, criteria, X, y, return_logits=False):\n    model.eval()\n    \n    logits_1, logits_2, logits_t = model(X)\n    \n    if return_logits:\n        return logits_1, logits_2, logits_t    \n    \n    ce_loss_1 = criteria['ce'](logits_1, y).item()\n    ce_loss_2 = criteria['ce'](logits_2, y).item()\n    ce_loss_t = criteria['ce'](logits_t, y).item()\n    total_loss = ce_loss_1 + ce_loss_2 + ce_loss_t\n\n    _, predicted_labels_1 = logits_1.max(1)\n    _, predicted_labels_2 = logits_2.max(1)\n    _, predicted_labels_t = logits_t.max(1)\n    acc_f1 = (predicted_labels_1 == y).to(torch.float).mean().item()\n    acc_f2 = (predicted_labels_2 == y).to(torch.float).mean().item()\n    acc_ft = (predicted_labels_t == y).to(torch.float).mean().item()\n    return total_loss, ce_loss_1, ce_loss_2, ce_loss_t, acc_f1, acc_f2, acc_ft\n\ndef judge_func(data, pred_1, pred_2, gt, pool_size, threshold=0.9):\n    data = data[:pool_size]\n    pred_1 = pred_1[:pool_size, :]\n    pred_2 = pred_2[:pool_size, :]\n    gt = gt[:pool_size, :]\n\n    p1 = pred_1.softmax(dim=1)\n    p2 = pred_2.softmax(dim=1)\n\n    n_samples = pred_1.shape[0]\n    new_data = []\n    new_label = []\n    new_gt = []\n\n    for i in range(n_samples):\n        cand_data = data[i, :, :, :]\n        max_prob_1 = p1[i].max()\n        max_prob_2 = p2[i].max()\n        c1 = p1[i].argmax()\n        c2 = p2[i].argmax()\n        if c1 == c2:\n            if max(max_prob_1, max_prob_2) > threshold:\n                new_label.append(c1)\n                new_data.append(cand_data)\n                new_gt.append(gt[i].to(torch.long))\n    \n    tqdm.write(f'Number of labeled: {len(new_label)}')\n    tqdm.write(f'Pool size: {pool_size}')\n    return torch.stack(new_data, dim=0), torch.tensor(new_label, dtype=torch.long, device=DEVICE), torch.stack(new_gt, dim=0).view(-1)\n\ndef validate(model, optims, criteria, t, t_val_loader, best_ft_acc):\n    model.eval()\n    with torch.no_grad():\n        for (x, y) in t_val_loader:\n            x, y = x.to(DEVICE), y.to(DEVICE)\n            _, _, _, _, _, _, val_acc_ft = run_eval(model, criteria, x, y, return_logits=False)\n            if val_acc_ft > best_ft_acc:\n                best_ft_acc = val_acc_ft\n                ckpt = {'model_weights': model.state_dict(),\n                        'fn_optim': optims['fn'].state_dict(),\n                        'ft_optim': optims['ft'].state_dict(),}\n                torch.save(ckpt, os.path.join(EXP_DIR, f'best_model_iter_{t}.pt'))\n    return best_ft_acc, val_acc_ft\n","metadata":{"execution":{"iopub.status.busy":"2023-08-25T07:53:36.044868Z","iopub.execute_input":"2023-08-25T07:53:36.045218Z","iopub.status.idle":"2023-08-25T07:53:36.079111Z","shell.execute_reply.started":"2023-08-25T07:53:36.045171Z","shell.execute_reply":"2023-08-25T07:53:36.078128Z"},"jupyter":{"source_hidden":true},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Helper functions for visualization","metadata":{}},{"cell_type":"code","source":"def plot_exp_results(acc_ft_history, label_acc_history, n_samples_history, exp_no=0, show_result=False):\n    _, ax1 = plt.subplots(dpi=120)\n    ax1.plot(acc_ft_history, marker='o', label='Target Net Accuracy')\n    ax1.plot(label_acc_history, marker='o', label='Pseudo Label Accuracy')\n    ax1.set_xlabel('Number of Iterations')\n    ax1.set_ylabel('Accuracy')\n\n    ax2 = ax1.twinx()\n    ax2.plot(n_samples_history, marker='s', label='Number of Samples', color='purple')\n    ax2.set_ylabel('Number of Samples')\n\n    # Combine legends from ax1 and ax2\n    lines_1, labels_1 = ax1.get_legend_handles_labels()\n    lines_2, labels_2 = ax2.get_legend_handles_labels()\n    lines = lines_1 + lines_2\n    labels = labels_1 + labels_2\n    ax1.legend(lines, labels, loc='upper center', bbox_to_anchor=(0.5, -0.15), fancybox=True, shadow=True, ncol=3)\n\n    # Set the number of grid lines on both axes\n    ax1.yaxis.set_major_locator(plt.MaxNLocator(len(ax2.get_yticks())))\n    ax2.yaxis.set_major_locator(plt.MaxNLocator(len(ax1.get_yticks())))    \n\n    plt.title(f'MNIST -> USPS | Global Best: {np.max(acc_ft_history)*100:.1f} | No. Iters. vs. Overall Accuracy')\n    plt.savefig(f'{EXP_DIR}/MNIST-USPS_results_exp_{exp_no}.png', dpi=120, bbox_inches='tight')\n    if show_result:\n        plt.show()    \n    plt.close()\n\ndef umap_visualization(model, s_loader, t_loader, title, iter=-1, n_samples=500):\n    model.eval()\n    s_embeddings, t_embeddings = [], []\n    s_y, t_y = [], []\n\n    with torch.no_grad():\n        for loader, embeddings, y in [(s_loader, s_embeddings, s_y), (t_loader, t_embeddings, t_y)]:\n            for x, y_batch in loader:\n                x = x.to(DEVICE)\n                embedding = model.get_embeddings(x)\n                embeddings.append(embedding.cpu().numpy())\n                y.append(y_batch.numpy())\n\n    s_y = np.concatenate(s_y, axis=0)[:n_samples]\n    t_y = np.concatenate(t_y, axis=0)[:n_samples]\n    s_embeddings = np.concatenate(s_embeddings, axis=0)[:n_samples]\n    t_embeddings = np.concatenate(t_embeddings, axis=0)[:n_samples]\n\n    all_embeddings = np.concatenate([s_embeddings, t_embeddings], axis=0)\n    reducer = umap.UMAP()\n    umap_embeddings = reducer.fit_transform(all_embeddings)\n\n    plt.figure(dpi=120)\n    plt.scatter(umap_embeddings[:n_samples, 0], umap_embeddings[:n_samples, 1], label='MNIST', s=7.5, marker='o', c='#1f77b4', alpha=0.75)\n    plt.scatter(umap_embeddings[n_samples:, 0], umap_embeddings[n_samples:, 1], label='USPS', s=7.5, marker='o', c='#ff7f0e', alpha=0.75)\n\n    # Plot cluster centers for each group\n    def plot_cluster_centers(embeddings, color, marker):\n        cluster_centers = np.mean(embeddings, axis=0)\n        plt.scatter(cluster_centers[0], cluster_centers[1], s=50, c=color, marker=marker)\n\n    plot_cluster_centers(umap_embeddings[:n_samples], '#1f77b4', 'o')\n    plot_cluster_centers(umap_embeddings[n_samples:], '#ff7f0e', '^')\n\n    plt.legend()\n    plt.title(title)\n    plt.savefig(f'{EXP_DIR}/umap/umap_iter_{iter}.png')\n    plt.close()","metadata":{"execution":{"iopub.status.busy":"2023-08-25T07:53:36.082216Z","iopub.execute_input":"2023-08-25T07:53:36.082567Z","iopub.status.idle":"2023-08-25T07:53:36.103121Z","shell.execute_reply.started":"2023-08-25T07:53:36.082540Z","shell.execute_reply":"2023-08-25T07:53:36.102174Z"},"jupyter":{"source_hidden":true},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Training loop\n![](https://i.imgur.com/3kWxnOO.png)","metadata":{}},{"cell_type":"code","source":"acc_ft_history = []\nlabel_acc_history = []\nn_samples_history = []\nfor t in range(N_ITERS):\n    print(f'Phase {t}')\n    if t == 0:\n        s_train_dataset = datasets.MNIST(SOURCE_DIR, train=True, transform=S_TRANSFORM, download=True)\n        s_train_iterator = cycle(DataLoader(s_train_dataset, batch_size=BATCH_SIZE, shuffle=True))\n        t_dataset = USPS(TARGET_DIR, train=True, transform=T_TRANSFORM, download=True)\n        \n        # Create indices for splitting\n        indices = list(range(len(t_dataset)))\n        train_indices = indices[:-500]\n        val_indices = indices[-500:]\n\n        # Create training and validation subsets\n        t_train_dataset = Subset(t_dataset, train_indices)\n        t_val_dataset = Subset(t_dataset, val_indices)        \n        t_val_loader = DataLoader(t_val_dataset, batch_size=BATCH_SIZE, shuffle=False)\n    else:\n        # Combine the source dataset and psuedo-labeled target domain dataset\n        tqdm.write('Preparing dataloaders...')\n        s_train_x, s_train_y = extract_datasets(s_train_dataset)\n        s_train = torch.cat([s_train_x.to(DEVICE), new_data], dim=0)\n        s_label = torch.cat([s_train_y.to(DEVICE), new_label], dim=0)\n        \n        combined_dataset = CustomImageDataset(s_train, s_label)\n        should_drop_last = len(combined_dataset) % int(BATCH_SIZE / 2) == 1\n        combined_iterator = cycle(DataLoader(combined_dataset, batch_size=int(BATCH_SIZE / 2), shuffle=True, drop_last=should_drop_last))\n        \n        new_dataset = CustomImageDataset(new_data, new_label)\n        should_drop_last = len(new_dataset) % BATCH_SIZE == 1\n        new_iterator = cycle(DataLoader(new_dataset, batch_size=BATCH_SIZE, shuffle=True, drop_last=should_drop_last))\n\n    # Training loop\n    best_ft_acc = 0\n    for i in tqdm(range(N_STEPS)):\n        if t == 0:\n            model.train()\n            x0, y0 = next(s_train_iterator)\n            x0, y0 = x0.to(DEVICE), y0.to(DEVICE)\n            batch_loss, wc_loss, ce_loss_1, ce_loss_2, _, acc_f1, acc_f2, _ = run_pretrain(model, optims, criteria, x0, y0)\n\n            if i % 500 == 0:\n                best_ft_acc, val_acc_ft = validate(model, optims, criteria, t, t_val_loader, best_ft_acc)\n                tqdm.write(f'Batch Loss: {batch_loss:.3f}, WC Loss: {wc_loss:.3f}, CE Loss 1: {ce_loss_1:.3f}, CE Loss 2: {ce_loss_2:.3f}, Acc F1: {acc_f1*100:.1f}, Acc F2: {acc_f2*100:.1f}, Val Acc Ft: {val_acc_ft*100:.1f}')\n\n        if t >= 1:\n            model.train()\n            x0, y0 = next(combined_iterator)\n            x0, y0 = x0.to(DEVICE), y0.to(DEVICE)\n            batch_loss, wc_loss, ce_loss_1, ce_loss_2, acc_f1, acc_f2 = run_retrain(model, optims, criteria, x0, y0, mode='mixed')\n\n            x1, y1 = next(new_iterator)\n            x1, y1 = x1.to(DEVICE), y1.to(DEVICE)\n            acc_ft = run_retrain(model, optims, criteria, x1, y1, mode='target')\n\n            if i % 500 == 0:\n                best_ft_acc, val_acc_ft = validate(model, optims, criteria, t, t_val_loader, best_ft_acc)\n                tqdm.write(f'Batch Loss: {batch_loss:.3f}, WC Loss: {wc_loss:.3f}, CE Loss 1: {ce_loss_1:.3f}, CE Loss 2: {ce_loss_2:.3f}, Acc F1: {acc_f1*100:.1f}, Acc F2: {acc_f2*100:.1f}, Acc Ft: {acc_ft*100:.1f}, Val Acc Ft: {val_acc_ft*100:.1f}')\n        \n    # Assign pseudo labels\n    tqdm.write('Assigning Pseudo-labels...')\n    candidates_x = []\n    pred_1_stack = torch.zeros((0, N_CLASSES), device=DEVICE)\n    pred_2_stack = torch.zeros((0, N_CLASSES), device=DEVICE)\n    gt_stack = torch.zeros((0, 1), device=DEVICE)\n    t_train_iterator = iter(DataLoader(t_train_dataset, batch_size=BATCH_SIZE, shuffle=True))\n    stack_num = int(min(len(t_train_dataset) / BATCH_SIZE, 100 * (t + 1)))\n    for _ in tqdm(range(stack_num)):\n        model.eval()\n        X, y = next(t_train_iterator)\n        X, y = X.to(DEVICE), y.to(DEVICE)\n        candidates_x.append(X)\n\n        with torch.no_grad():\n            pred_1, pred_2, _ = run_eval(model, criteria, X, y, return_logits=True)\n\n        pred_1_stack = torch.vstack([pred_1_stack, pred_1])\n        pred_2_stack = torch.vstack([pred_2_stack, pred_2])\n        gt_stack = torch.cat([gt_stack, y.unsqueeze(1)], dim=0)\n\n    candidates_x = torch.cat(candidates_x, dim=0)\n    if t == 0:\n        pool_size = max(int((t + 1) / 20 * pred_1_stack.shape[0]), 1000)\n    else:\n        pool_size = min(max(int((t + 1) / 20 * pred_1_stack.shape[0]), 1000), 40000)\n    new_data, new_label, new_gt = judge_func(candidates_x, pred_1_stack, pred_2_stack, gt_stack, pool_size, threshold=THRESHOLD)\n    label_acc = (new_gt.view(-1) == new_label.view(-1)).to(torch.float).mean().item()\n    tqdm.write(f'Pseudo-label accuracy: {label_acc*100:.1f}')\n\n    # Evaluation\n    tqdm.write('Evaluating...')\n    s_test_dataset = datasets.MNIST(SOURCE_DIR, train=False, transform=S_TRANSFORM, download=True)\n    t_test_dataset = USPS(TARGET_DIR, train=False, transform=T_TRANSFORM, download=True)\n    s_test_loader = DataLoader(s_test_dataset, batch_size=BATCH_SIZE, shuffle=False)\n    t_test_loader = DataLoader(t_test_dataset, batch_size=BATCH_SIZE, shuffle=False)\n    s_test_iterator = cycle(s_test_loader)\n    t_test_iterator = iter(t_test_loader)\n\n    total_source = 0\n    total_target = 0\n    total_acc_1 = 0\n    total_acc_2 = 0\n    s_size = 0\n    t_size = 0\n\n    n_iter = int(len(t_test_loader.dataset) / BATCH_SIZE) + 1\n    for _ in tqdm(range(n_iter)):\n        x0, y0 = next(s_test_iterator)\n        x1, y1 = next(t_test_iterator)\n        x0, y0 = x0.to(DEVICE), y0.to(DEVICE)\n        x1, y1 = x1.to(DEVICE), y1.to(DEVICE)\n\n        with torch.no_grad():\n            _, _, _, _, s_acc, _, _ = run_eval(model, criteria, x0, y0, return_logits=False)\n            _, _, _, _, t_acc_f1, t_acc_f2, t_acc_ft = run_eval(model, criteria, x1, y1, return_logits=False)\n\n        total_source += s_acc * x0.shape[0]\n        total_target += t_acc_ft * x1.shape[0]\n        total_acc_1 += t_acc_f1 * x1.shape[0]\n        total_acc_2 += t_acc_f2 * x1.shape[0]\n        s_size += x0.shape[0]\n        t_size += x1.shape[0]\n\n    total_acc_f1 = total_acc_1 / t_size\n    total_acc_f2 = total_acc_2 / t_size\n    total_acc_ft = total_target / t_size\n    total_acc_s = total_source / s_size\n    tqdm.write(f'Target Ft Acc: {total_acc_ft*100:.1f}, Target F1 Acc: {total_acc_f1*100:.1f}, Target F2 Acc: {total_acc_f2*100:.1f}, Source F1 Acc: {total_acc_s*100:.1f}')\n\n    acc_ft_history.append(total_acc_ft)\n    label_acc_history.append(label_acc)\n    n_samples_history.append(new_label.numel())        \n\n    plot_exp_results(acc_ft_history, label_acc_history, n_samples_history)\n    umap_visualization(model, s_test_loader, t_test_loader, iter=t, n_samples=1000, title=f'UMAP Visualization at Iteration {t}')\n\nplot_exp_results(acc_ft_history, label_acc_history, n_samples_history, show_result=True)\ntotal_acc_s = np.max(total_acc_s)\ntotal_acc_ft = np.max(total_acc_ft)\ntotal_acc_f1 = np.max(total_acc_f1)\ntotal_acc_f2 = np.max(total_acc_f2)","metadata":{"execution":{"iopub.status.busy":"2023-08-25T07:53:36.104682Z","iopub.execute_input":"2023-08-25T07:53:36.105047Z","iopub.status.idle":"2023-08-25T08:13:18.910665Z","shell.execute_reply.started":"2023-08-25T07:53:36.105014Z","shell.execute_reply":"2023-08-25T08:13:18.909678Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Export experimental results","metadata":{}},{"cell_type":"code","source":"import subprocess\nfrom IPython.display import FileLink, display\n\ndef download_file(path, download_file_name):\n    os.chdir('/kaggle/working/')\n    zip_name = f\"/kaggle/working/{download_file_name}.zip\"\n    command = f\"zip {zip_name} {path} -r\"\n    result = subprocess.run(command, shell=True, capture_output=True, text=True)\n    if result.returncode != 0:\n        print(\"Unable to run zip command!\")\n        print(result.stderr)\n        return\n    display(FileLink(f'{download_file_name}.zip'))\n    \ndownload_file(EXP_DIR, 'exps')","metadata":{"execution":{"iopub.status.busy":"2023-08-25T08:13:18.911808Z","iopub.execute_input":"2023-08-25T08:13:18.912904Z","iopub.status.idle":"2023-08-25T08:13:22.767938Z","shell.execute_reply.started":"2023-08-25T08:13:18.912862Z","shell.execute_reply":"2023-08-25T08:13:22.766999Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}