update code.

.gitignore

@@ -0,0 +1,10 @@
+/.idea
+/.vscode
+**/checkpoint*/
+**/__pycache__/
+**/generated*/
+**/wandb/
+**/full_model.pth
+/rubbish
+**/*cache*
+/workspace/classinput/Qwen25llm/

README.md

@@ -1,3 +1,79 @@
----
-license: unknown
----
+# Recurrent Parameter Generation
+The official repository of paper [Recurrent Diffusion for Large-Scale Parameter Generation]().
+
+
+##  Introduction
+Parameter generation has long struggled to scale, significantly limiting its applications. 
+In this study, we introduce Recurrent diffusion for large-scale Parameter Generation, or RPG, 
+which models large-scale parameter generation through a recurrent diffusion process. 
+We divide the trained parameters into non-overlapping parts and propose a recurrent model to learn their relationships. 
+The outputs of this recurrent model, serving as conditions, are then input into a diffusion model to generate neural network parameters. 
+Utilizing only a single GPU, our method can generate parameters for popular vision and language models, such as ConvNeXt-L and LoRA parameters for LLaMA-7B. 
+Across various architectures and tasks, the generated parameters consistently achieve comparable performance to those of trained networks. 
+Additionally, our approach demonstrates potential in generating models capable of handling unseen tasks, 
+indicating that recurrent diffusion greatly enhances the practicality of parameter generation.
+
+
+
+
+
+
+
+
+## Environment
+Before you get started, you need to set up a conda environment first.
+1. Create your conda environment.
+```shell
+conda create -n rpg python=3.11
+conda activate rpg
+conda install pytorch==2.3.1 torchvision==0.18.1 torchaudio==2.3.1 pytorch-cuda=12.1 -c pytorch -c nvidia
+```
+2. Install mamba-ssm. (You may run into compilation issues, refer to the [official mamba-ssm repository](https://github.com/state-spaces/mamba) for details.)
+```shell
+pip install mamba-ssm[causal-conv1d]
+pip install causal-conv1d
+```
+3. Install other dependencies for this repository.
+```shell
+git lfs install
+git clone https://huggingface.co/MTDoven/Recurrent-Parameter-Generation
+cd Recurrent-Parameter-Generation
+pip install -r requirements.txt
+```
+
+
+
+
+## Quick Start
+1. Modify your config file.
+```shell
+# Set up your configs interactively.
+python ./workspace/set_configs.py
+```
+
+2. Download checkpoint datasets.
+```shell
+
+# Download the ViTTiny1022 dataset to /path/to/your/download/ViTTiny1022
+mv /path/to/your/download/ViTTiny1022/* ./dataset/condition_classinput_vittiny/
+```
+
+3. Try to generate with RPG model.
+```shell
+cd ./workspace
+CUDA_VISIBLE_DEVICES=0 python ./classinput/launch.py
+# CUDA_VISIBLE_DEVICES=<GPU_index> python ./classinput/launch.py
+```
+
+You can get more information from [Github](https://github.com/NUS-HPC-AI-Lab/Recurrent-Parameter-Generation).
+
+
+
+
+## Acknowledgment
+coming soon...
+
+
+## Citation
+coming soon...
+

dataset/__init__.py

@@ -0,0 +1 @@
+from .register import *

dataset/cifar100_resnet18bn/model.py

@@ -0,0 +1,27 @@
+import torch.nn as nn
+import torch
+import timm
+import os
+
+
+def Model():
+    model = timm.create_model("resnet18", pretrained=True)
+    model.fc = nn.Linear(512, 100)
+    if os.path.exists(os.path.join(os.path.dirname(__file__), "full_model.pth")):
+        model.load_state_dict(torch.load(os.path.join(os.path.dirname(__file__), "full_model.pth"), map_location="cpu"))
+        for k, v in model.named_parameters():
+            if k in ["layer4.1.bn1.weight", "layer4.1.bn1.bias", "layer4.1.bn2.weight", "layer4.1.bn2.bias"]:
+                v.requires_grad = True
+            else:  # requires_grad = False
+                v.requires_grad = False
+    return model, model.fc
+
+
+if __name__ == "__main__":
+    model, _ = Model()
+    print(model)
+    num_param = 0
+    for k, v in model.named_parameters():
+        num_param += v.numel()
+        print(k)
+    print("num_param:", num_param)

dataset/cifar100_resnet18bn/prepare.py

@@ -0,0 +1,192 @@
+# set global seed
+import random
+import numpy as np
+import torch
+seed = SEED = 20
+torch.manual_seed(seed)
+torch.cuda.manual_seed(seed)
+torch.cuda.manual_seed_all(seed)
+torch.backends.cudnn.deterministic = True
+torch.backends.cudnn.benchmark = True
+np.random.seed(seed)
+random.seed(seed)
+
+
+try:  # relative import
+    from model import Model
+except ImportError:
+    from .model import Model
+
+# import
+import torch.nn as nn
+from torch import optim
+from torch.optim import lr_scheduler
+from torch.utils.data import DataLoader
+import torchvision.transforms as transforms
+from torchvision.datasets import CIFAR100 as Dataset
+from tqdm.auto import tqdm
+import os
+import warnings
+warnings.filterwarnings("ignore", category=UserWarning)
+
+# load additional config
+import json
+config_file = os.path.join(os.path.dirname(os.path.dirname(__file__)), "config.json")
+with open(config_file, "r") as f:
+    additional_config = json.load(f)
+
+
+
+
+# config
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+config = {
+    "dataset_root": "from_additional_config",
+    "batch_size": 500 if __name__ == "__main__" else 200,
+    "num_workers": 32,
+    "learning_rate": 0.0005,
+    "weight_decay": 0.000005,
+    "epochs": 200,
+    "save_learning_rate": 0.0,
+    "total_save_number": 1,
+    "tag": os.path.basename(os.path.dirname(__file__)),
+}
+config.update(additional_config)
+
+
+
+
+# Data
+dataset = Dataset(
+    root=config["dataset_root"],
+    download=True,
+    train=True,
+    transform=transforms.Compose([
+        transforms.Resize(80),
+        transforms.RandomHorizontalFlip(),
+        transforms.AutoAugment(policy=transforms.AutoAugmentPolicy("cifar10")),
+        transforms.ToTensor(),
+        transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
+    ])
+)
+train_loader = DataLoader(
+    dataset=dataset,
+    batch_size=config["batch_size"],
+    num_workers=config["num_workers"],
+    shuffle=True,
+    drop_last=True,
+    pin_memory=True,
+)
+test_loader = DataLoader(
+    dataset=Dataset(
+        root=config["dataset_root"],
+        download=True,
+        train=False,
+        transform=transforms.Compose([
+            transforms.Resize(80),
+            transforms.ToTensor(),
+            transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
+        ])),
+    batch_size=config["batch_size"],
+    num_workers=config["num_workers"],
+    shuffle=False,
+    pin_memory=True,
+)
+
+# Model
+model, head = Model()
+model = model.to(device)
+criterion = nn.CrossEntropyLoss()
+pre_optimizer = optim.AdamW(
+    head.parameters(),
+    lr=0.001,
+    weight_decay=config["weight_decay"],
+)
+optimizer = optim.AdamW(
+    model.parameters(),
+    lr=config["learning_rate"],
+    weight_decay=config["weight_decay"],
+)
+scheduler = lr_scheduler.CosineAnnealingLR(
+    optimizer,
+    T_max=config["epochs"],
+    eta_min=config["save_learning_rate"],
+)
+
+
+
+
+# Training
+def train(model=model, optimizer=optimizer, scheduler=scheduler):
+    model.train()
+    for batch_idx, (inputs, targets) in tqdm(enumerate(train_loader),
+                                             total=len(dataset) // config["batch_size"]):
+        inputs, targets = inputs.to(device), targets.to(device)
+        optimizer.zero_grad()
+        with torch.cuda.amp.autocast(enabled=False, dtype=torch.bfloat16):
+            outputs = model(inputs)
+            loss = criterion(outputs, targets)
+        loss.backward()
+        optimizer.step()
+    if scheduler is not None:
+        scheduler.step()
+
+# test
+@torch.no_grad()
+def test(model=model):
+    model.eval()
+    all_targets = []
+    all_predicts = []
+    test_loss = 0
+    correct = 0
+    total = 0
+    for batch_idx, (inputs, targets) in tqdm(enumerate(test_loader),
+                                             total=len(test_loader.dataset) // config["batch_size"]):
+        inputs, targets = inputs.to(device), targets.to(device)
+        with torch.cuda.amp.autocast(enabled=False, dtype=torch.bfloat16):
+            outputs = model(inputs)
+            loss = criterion(outputs, targets)
+        # to logging losses
+        all_targets.extend(targets.flatten().tolist())
+        test_loss += loss.item()
+        _, predicts = outputs.max(1)
+        all_predicts.extend(predicts.flatten().tolist())
+        total += targets.size(0)
+        correct += predicts.eq(targets).sum().item()
+    loss = test_loss / (batch_idx + 1)
+    acc = correct / total
+    print(f"Loss: {loss:.4f} | Acc: {acc:.4f}")
+    model.train()
+    return loss, acc, all_targets, all_predicts
+
+# save train
+def save_train(model=model, optimizer=optimizer):
+    model.train()
+    for batch_idx, (inputs, targets) in enumerate(train_loader):
+        inputs, targets = inputs.to(device), targets.to(device)
+        optimizer.zero_grad()
+        with torch.cuda.amp.autocast(enabled=False, dtype=torch.bfloat16):
+            outputs = model(inputs)
+            loss = criterion(outputs, targets)
+        loss.backward()
+        optimizer.step()
+    # Save checkpoint
+    _, acc, _, _ = test(model=model)
+    if not os.path.isdir('checkpoint'):
+        os.mkdir('checkpoint')
+    save_state = {key: value.cpu().to(torch.float32) for key, value in model.state_dict().items()}
+    torch.save(save_state, f"full_model.pth")
+    print("save:", f"full_model.pth")
+
+
+
+
+# main
+if __name__ == '__main__':
+    test(model=model)
+    train(model=model, optimizer=pre_optimizer, scheduler=scheduler)
+    train(model=model, optimizer=pre_optimizer, scheduler=scheduler)
+    for epoch in range(config["epochs"]):
+        train(model=model, optimizer=optimizer, scheduler=scheduler)
+        test(model=model)
+    save_train(model=model, optimizer=optimizer)

dataset/cifar100_resnet18bn/test.py

@@ -0,0 +1,28 @@
+import os
+import sys
+if __name__ == "__main__":
+    from train import *
+else:  # relative import
+    from .train import *
+
+
+
+
+try:
+    test_item = sys.argv[1]
+except IndexError:
+    assert __name__ == "__main__"
+    test_item = "./checkpoint"
+test_items = []
+if os.path.isdir(test_item):
+    for item in os.listdir(test_item):
+        item = os.path.join(test_item, item)
+        test_items.append(item)
+elif os.path.isfile(test_item):
+    test_items.append(test_item)
+
+
+for item in test_items:
+    state = torch.load(item, map_location="cpu")
+    model.load_state_dict({key: value.to(torch.float32).to(device) for key, value in state.items()}, strict=False)
+    loss, acc, all_targets, all_predicts = test(model=model)

dataset/cifar100_resnet18bn/train.py

@@ -0,0 +1,195 @@
+# set global seed
+import random
+import numpy as np
+import torch
+seed = SEED = 20
+torch.manual_seed(seed)
+torch.cuda.manual_seed(seed)
+torch.cuda.manual_seed_all(seed)
+torch.backends.cudnn.deterministic = True
+torch.backends.cudnn.benchmark = True
+np.random.seed(seed)
+random.seed(seed)
+
+
+try:  # relative import
+    from model import Model
+except ImportError:
+    from .model import Model
+
+# import
+import torch.nn as nn
+from torch import optim
+from torch.optim import lr_scheduler
+from torch.utils.data import DataLoader
+import torchvision.transforms as transforms
+from torchvision.datasets import CIFAR100 as Dataset
+from tqdm.auto import tqdm
+import os
+import warnings
+warnings.filterwarnings("ignore", category=UserWarning)
+
+# load additional config
+import json
+config_file = os.path.join(os.path.dirname(os.path.dirname(__file__)), "config.json")
+with open(config_file, "r") as f:
+    additional_config = json.load(f)
+
+
+
+
+# config
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+config = {
+    "dataset_root": "from_additional_config",
+    "batch_size": 100 if __name__ == "__main__" else 200,
+    "num_workers": 4,
+    "learning_rate": 0.01,
+    "weight_decay": 5e-6,
+    "epochs": 1,
+    "save_learning_rate": 0.01,
+    "total_save_number": 200,
+    "tag": os.path.basename(os.path.dirname(__file__)),
+}
+config.update(additional_config)
+
+
+
+
+# Data
+dataset = Dataset(
+    root=config["dataset_root"],
+    download=True,
+    train=True,
+    transform=transforms.Compose([
+        transforms.Resize(80),
+        transforms.RandomHorizontalFlip(),
+        transforms.RandAugment(),
+        transforms.ToTensor(),
+        transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
+    ])
+)
+train_loader = DataLoader(
+    dataset=dataset,
+    batch_size=config["batch_size"],
+    num_workers=config["num_workers"],
+    shuffle=True,
+    drop_last=True,
+    pin_memory=True,
+    persistent_workers=False,
+)
+test_loader = DataLoader(
+    dataset=Dataset(
+        root=config["dataset_root"],
+        download=True,
+        train=False,
+        transform=transforms.Compose([
+            transforms.Resize(80),
+            transforms.CenterCrop(80),
+            transforms.ToTensor(),
+            transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
+        ])),
+    batch_size=config["batch_size"],
+    num_workers=config["num_workers"],
+    shuffle=False,
+    pin_memory=True,
+    persistent_workers=False,
+    pin_memory_device="cuda",
+)
+
+# Model
+model, head = Model()
+model = model.to(device)
+criterion = nn.CrossEntropyLoss()
+optimizer = optim.AdamW(
+    model.parameters(),
+    lr=config["learning_rate"],
+    weight_decay=config["weight_decay"],
+)
+scheduler = lr_scheduler.CosineAnnealingLR(
+    optimizer,
+    T_max=config["epochs"],
+    eta_min=config["save_learning_rate"],
+)
+
+
+
+
+# Training
+def train(model=model, optimizer=optimizer, scheduler=scheduler):
+    model.train()
+    for batch_idx, (inputs, targets) in tqdm(enumerate(train_loader),
+                                             total=len(dataset) // config["batch_size"]):
+        inputs, targets = inputs.to(device), targets.to(device)
+        optimizer.zero_grad()
+        with torch.cuda.amp.autocast(enabled=False, dtype=torch.bfloat16):
+            outputs = model(inputs)
+            loss = criterion(outputs, targets)
+        loss.backward()
+        optimizer.step()
+    if scheduler is not None:
+        scheduler.step()
+
+# test
+@torch.no_grad()
+def test(model=model):
+    model.eval()
+    all_targets = []
+    all_predicts = []
+    test_loss = 0
+    correct = 0
+    total = 0
+    for batch_idx, (inputs, targets) in tqdm(enumerate(test_loader),
+                                             total=len(test_loader.dataset) // config["batch_size"]):
+        inputs, targets = inputs.to(device), targets.to(device)
+        with torch.cuda.amp.autocast(enabled=False, dtype=torch.bfloat16):
+            outputs = model(inputs)
+            loss = criterion(outputs, targets)
+        # to logging losses
+        all_targets.extend(targets.flatten().tolist())
+        test_loss += loss.item()
+        _, predicts = outputs.max(1)
+        all_predicts.extend(predicts.flatten().tolist())
+        total += targets.size(0)
+        correct += predicts.eq(targets).sum().item()
+    loss = test_loss / (batch_idx + 1)
+    acc = correct / total
+    print(f"Loss: {loss:.4f} | Acc: {acc:.4f}\n")
+    model.train()
+    return loss, acc, all_targets, all_predicts
+
+# save train
+def save_train(model=model, optimizer=optimizer):
+    model.train()
+    saved_number = 0
+    for batch_idx, (inputs, targets) in enumerate(train_loader):
+        inputs, targets = inputs.to(device), targets.to(device)
+        optimizer.zero_grad()
+        with torch.cuda.amp.autocast(enabled=False, dtype=torch.bfloat16):
+            outputs = model(inputs)
+            loss = criterion(outputs, targets)
+        loss.backward()
+        optimizer.step()
+        # Save checkpoint
+        if batch_idx % (len(dataset) // train_loader.batch_size // config["total_save_number"]) == 0:
+            _, acc, _, _ = test(model=model)
+            if not os.path.isdir('checkpoint'):
+                os.mkdir('checkpoint')
+            save_state = {key: value.cpu().to(torch.float32) for key, value in model.state_dict().items() \
+                          if key in ["layer4.1.bn1.weight", "layer4.1.bn1.bias", "layer4.1.bn2.weight", "layer4.1.bn2.bias"]}
+            torch.save(save_state, f"checkpoint/{str(batch_idx).zfill(4)}_acc{acc:.4f}_seed{seed:04d}_{config['tag']}.pth")
+            print("save:", f"checkpoint/{str(batch_idx).zfill(4)}_acc{acc:.4f}_seed{seed:04d}_{config['tag']}.pth")
+            saved_number += 1
+            if saved_number >= config["total_save_number"]:
+                break
+
+
+
+
+# main
+if __name__ == '__main__':
+    test(model=model)
+    for epoch in range(config["epochs"]):
+        train(model=model, optimizer=optimizer, scheduler=scheduler)
+        test(model=model)
+    save_train(model=model, optimizer=optimizer)

dataset/cifar10_cnnmedium/model.py

@@ -0,0 +1,48 @@
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+import timm
+
+
+class CNNMedium(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.module = nn.Sequential(
+            nn.Conv2d(3, 16, 3),
+            nn.MaxPool2d(2, 2),
+            nn.LeakyReLU(),
+            nn.Conv2d(16, 32, 3),
+            nn.MaxPool2d(2, 2),
+            nn.LeakyReLU(),
+            nn.Conv2d(32, 15, 3),
+            nn.MaxPool2d(2, 2),
+            nn.LeakyReLU(),
+            nn.Flatten(start_dim=1),
+        )
+        self.head = nn.Sequential(
+            nn.Linear(60, 20),
+            nn.LeakyReLU(),
+            nn.Linear(20, 10),
+        )
+
+    def forward(self, x):
+        x = self.module(x)
+        x = self.head(x)
+        return x
+
+
+def Model():
+    model = CNNMedium()
+    return model, model.head
+
+
+if __name__ == "__main__":
+    model, _ = Model()
+    x = torch.ones([4, 3, 32, 32])
+    y = model(x)
+    print(y.shape)
+    print(model)
+    num_param = 0
+    for v in model.parameters():
+        num_param += v.numel()
+    print("num_param:", num_param)

dataset/cifar10_cnnmedium/test.py

@@ -0,0 +1,28 @@
+import os
+import sys
+if __name__ == "__main__":
+    from train import *
+else:  # relative import
+    from .train import *
+
+
+
+
+try:
+    test_item = sys.argv[1]
+except IndexError:
+    assert __name__ == "__main__"
+    test_item = "./checkpoint"
+test_items = []
+if os.path.isdir(test_item):
+    for item in os.listdir(test_item):
+        item = os.path.join(test_item, item)
+        test_items.append(item)
+elif os.path.isfile(test_item):
+    test_items.append(test_item)
+
+
+for item in test_items:
+    state = torch.load(item, map_location="cpu")
+    model.load_state_dict({key: value.to(torch.float32).to(device) for key, value in state.items()})
+    loss, acc, all_targets, all_predicts = test(model=model)

dataset/cifar10_cnnmedium/train.py

@@ -0,0 +1,192 @@
+# set global seed
+import random
+import numpy as np
+import torch
+seed = SEED = 20
+torch.manual_seed(seed)
+torch.cuda.manual_seed(seed)
+torch.cuda.manual_seed_all(seed)
+torch.backends.cudnn.deterministic = True
+torch.backends.cudnn.benchmark = True
+np.random.seed(seed)
+random.seed(seed)
+
+
+try:  # relative import
+    from model import Model
+except ImportError:
+    from .model import Model
+
+# import
+import torch.nn as nn
+from torch import optim
+from torch.optim import lr_scheduler
+from torch.utils.data import DataLoader
+import torchvision.transforms as transforms
+from torchvision.datasets import CIFAR10 as Dataset
+from tqdm.auto import tqdm
+import os
+import warnings
+warnings.filterwarnings("ignore", category=UserWarning)
+
+# load additional config
+import json
+config_file = os.path.join(os.path.dirname(os.path.dirname(__file__)), "config.json")
+with open(config_file, "r") as f:
+    additional_config = json.load(f)
+
+
+
+
+# config
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+config = {
+    "dataset_root": "from_additional_config",
+    "batch_size": 500 if __name__ == "__main__" else 200,
+    "num_workers": 32,
+    "learning_rate": 1e-2,
+    "weight_decay": 0.00666,
+    "epochs": 50,
+    "save_learning_rate": 1e-5,
+    "total_save_number": 50,
+    "tag": os.path.basename(os.path.dirname(__file__)),
+}
+config.update(additional_config)
+
+
+
+
+# Data
+dataset = Dataset(
+    root=config["dataset_root"],
+    download=True,
+    train=True,
+    transform=transforms.Compose([
+        transforms.Resize(32),
+        transforms.RandomCrop(32),
+        transforms.RandomHorizontalFlip(),
+        transforms.AutoAugment(policy=transforms.AutoAugmentPolicy("cifar10")),
+        transforms.ToTensor(),
+        transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
+    ])
+)
+train_loader = DataLoader(
+    dataset=dataset,
+    batch_size=config["batch_size"],
+    num_workers=config["num_workers"],
+    shuffle=True,
+    drop_last=True,
+    pin_memory=True,
+    persistent_workers=True,
+)
+test_loader = DataLoader(
+    dataset=Dataset(
+        root=config["dataset_root"],
+        download=True,
+        train=False,
+        transform=transforms.Compose([
+            transforms.Resize(32),
+            transforms.CenterCrop(32),
+            transforms.ToTensor(),
+            transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
+        ])),
+    batch_size=config["batch_size"],
+    num_workers=config["num_workers"],
+    shuffle=False,
+    pin_memory=True,
+    persistent_workers=True,
+    pin_memory_device="cuda",
+)
+
+# Model
+model, head = Model()
+model = model.to(device)
+criterion = nn.CrossEntropyLoss()
+optimizer = optim.SGD(
+    model.parameters(),
+    lr=config["learning_rate"],
+    weight_decay=config["weight_decay"],
+    momentum=0.9,
+)
+scheduler = lr_scheduler.CosineAnnealingLR(
+    optimizer,
+    T_max=config["epochs"],
+    eta_min=config["save_learning_rate"],
+)
+
+
+
+
+# Training
+def train(model=model, optimizer=optimizer, scheduler=scheduler):
+    model.train()
+    for batch_idx, (inputs, targets) in tqdm(enumerate(train_loader),
+                                             total=len(dataset) // config["batch_size"]):
+        inputs, targets = inputs.to(device), targets.to(device)
+        optimizer.zero_grad()
+        with torch.cuda.amp.autocast(enabled=True, dtype=torch.bfloat16):
+            outputs = model(inputs)
+            loss = criterion(outputs, targets)
+        loss.backward()
+        optimizer.step()
+    if scheduler is not None:
+        scheduler.step()
+
+# test
+@torch.no_grad()
+def test(model=model):
+    model.eval()
+    all_targets = []
+    all_predicts = []
+    test_loss = 0
+    correct = 0
+    total = 0
+    for batch_idx, (inputs, targets) in tqdm(enumerate(test_loader),
+                                             total=len(test_loader.dataset) // config["batch_size"]):
+        inputs, targets = inputs.to(device), targets.to(device)
+        with torch.cuda.amp.autocast(enabled=False, dtype=torch.bfloat16):
+            outputs = model(inputs)
+            loss = criterion(outputs, targets)
+        # to logging losses
+        all_targets.extend(targets.flatten().tolist())
+        test_loss += loss.item()
+        _, predicts = outputs.max(1)
+        all_predicts.extend(predicts.flatten().tolist())
+        total += targets.size(0)
+        correct += predicts.eq(targets).sum().item()
+    loss = test_loss / (batch_idx + 1)
+    acc = correct / total
+    print(f"Loss: {loss:.4f} | Acc: {acc:.4f}\n")
+    model.train()
+    return loss, acc, all_targets, all_predicts
+
+# save train
+def save_train(model=model, optimizer=optimizer):
+    model.train()
+    for batch_idx, (inputs, targets) in enumerate(train_loader):
+        inputs, targets = inputs.to(device), targets.to(device)
+        optimizer.zero_grad()
+        with torch.cuda.amp.autocast(enabled=True, dtype=torch.bfloat16):
+            outputs = model(inputs)
+            loss = criterion(outputs, targets)
+        loss.backward()
+        optimizer.step()
+        # Save checkpoint
+        if batch_idx % (len(dataset) // train_loader.batch_size // config["total_save_number"]) == 0:
+            _, acc, _, _ = test(model=model)
+            if not os.path.isdir('checkpoint'):
+                os.mkdir('checkpoint')
+            save_state = {key: value.cpu().to(torch.float32) for key, value in model.state_dict().items()}
+            torch.save(save_state, f"checkpoint/{str(batch_idx).zfill(4)}_acc{acc:.4f}_seed{seed:04d}_{config['tag']}.pth")
+            print("save:", f"checkpoint/{str(batch_idx).zfill(4)}_acc{acc:.4f}_seed{seed:04d}_{config['tag']}.pth")
+
+
+
+
+# main
+if __name__ == '__main__':
+    test(model=model)
+    for epoch in range(config["epochs"]):
+        train(model=model, optimizer=optimizer, scheduler=scheduler)
+        test(model=model)
+    save_train(model=model, optimizer=optimizer)

dataset/cifar10_cnnsmall/model.py

@@ -0,0 +1,48 @@
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+import timm
+
+
+class CNNSmall(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.module = nn.Sequential(
+            nn.Conv2d(3, 8, 5),
+            nn.MaxPool2d(2, 2),
+            nn.LeakyReLU(),
+            nn.Conv2d(8, 6, 5),
+            nn.MaxPool2d(2, 2),
+            nn.LeakyReLU(),
+            nn.Conv2d(6, 4, 2),
+            nn.LeakyReLU(),
+            nn.Flatten(start_dim=1),
+        )
+        self.head = nn.Sequential(
+            nn.Linear(36, 20),
+            nn.LeakyReLU(),
+            nn.Linear(20, 10),
+        )
+
+    def forward(self, x):
+        x = F.interpolate(x, (28, 28), mode='bilinear')
+        x = self.module(x)
+        x = self.head(x)
+        return x
+
+
+def Model():
+    model = CNNSmall()
+    return model, model.head
+
+
+if __name__ == "__main__":
+    model, _ = Model()
+    x = torch.ones([4, 3, 28, 28])
+    y = model(x)
+    print(y.shape)
+    print(model)
+    num_param = 0
+    for v in model.parameters():
+        num_param += v.numel()
+    print("num_param:", num_param)

dataset/cifar10_cnnsmall/test.py

@@ -0,0 +1,28 @@
+import os
+import sys
+if __name__ == "__main__":
+    from train import *
+else:  # relative import
+    from .train import *
+
+
+
+
+try:
+    test_item = sys.argv[1]
+except IndexError:
+    assert __name__ == "__main__"
+    test_item = "./checkpoint"
+test_items = []
+if os.path.isdir(test_item):
+    for item in os.listdir(test_item):
+        item = os.path.join(test_item, item)
+        test_items.append(item)
+elif os.path.isfile(test_item):
+    test_items.append(test_item)
+
+
+for item in test_items:
+    state = torch.load(item, map_location="cpu")
+    model.load_state_dict({key: value.to(torch.float32).to(device) for key, value in state.items()})
+    loss, acc, all_targets, all_predicts = test(model=model)

dataset/cifar10_cnnsmall/train.py

@@ -0,0 +1,192 @@
+# set global seed
+import random
+import numpy as np
+import torch
+seed = SEED = 20
+torch.manual_seed(seed)
+torch.cuda.manual_seed(seed)
+torch.cuda.manual_seed_all(seed)
+torch.backends.cudnn.deterministic = True
+torch.backends.cudnn.benchmark = True
+np.random.seed(seed)
+random.seed(seed)
+
+
+try:  # relative import
+    from model import Model
+except ImportError:
+    from .model import Model
+
+# import
+import torch.nn as nn
+from torch import optim
+from torch.optim import lr_scheduler
+from torch.utils.data import DataLoader
+import torchvision.transforms as transforms
+from torchvision.datasets import CIFAR10 as Dataset
+from tqdm.auto import tqdm
+import os
+import warnings
+warnings.filterwarnings("ignore", category=UserWarning)
+
+# load additional config
+import json
+config_file = os.path.join(os.path.dirname(os.path.dirname(__file__)), "config.json")
+with open(config_file, "r") as f:
+    additional_config = json.load(f)
+
+
+
+
+# config
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+config = {
+    "dataset_root": "from_additional_config",
+    "batch_size": 500 if __name__ == "__main__" else 200,
+    "num_workers": 32,
+    "learning_rate": 1e-2,
+    "weight_decay": 0.001,
+    "epochs": 50,
+    "save_learning_rate": 1e-5,
+    "total_save_number": 50,
+    "tag": os.path.basename(os.path.dirname(__file__)),
+}
+config.update(additional_config)
+
+
+
+
+# Data
+dataset = Dataset(
+    root=config["dataset_root"],
+    download=True,
+    train=True,
+    transform=transforms.Compose([
+        transforms.Resize(32),
+        transforms.RandomCrop(32),
+        transforms.RandomHorizontalFlip(),
+        transforms.AutoAugment(policy=transforms.AutoAugmentPolicy("cifar10")),
+        transforms.ToTensor(),
+        transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
+    ])
+)
+train_loader = DataLoader(
+    dataset=dataset,
+    batch_size=config["batch_size"],
+    num_workers=config["num_workers"],
+    shuffle=True,
+    drop_last=True,
+    pin_memory=True,
+    persistent_workers=True,
+)
+test_loader = DataLoader(
+    dataset=Dataset(
+        root=config["dataset_root"],
+        download=True,
+        train=False,
+        transform=transforms.Compose([
+            transforms.Resize(32),
+            transforms.CenterCrop(32),
+            transforms.ToTensor(),
+            transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
+        ])),
+    batch_size=config["batch_size"],
+    num_workers=config["num_workers"],
+    shuffle=False,
+    pin_memory=True,
+    persistent_workers=True,
+    pin_memory_device="cuda",
+)
+
+# Model
+model, head = Model()
+model = model.to(device)
+criterion = nn.CrossEntropyLoss()
+optimizer = optim.SGD(
+    model.parameters(),
+    lr=config["learning_rate"],
+    weight_decay=config["weight_decay"],
+    momentum=0.9,
+)
+scheduler = lr_scheduler.CosineAnnealingLR(
+    optimizer,
+    T_max=config["epochs"],
+    eta_min=config["save_learning_rate"],
+)
+
+
+
+
+# Training
+def train(model=model, optimizer=optimizer, scheduler=scheduler):
+    model.train()
+    for batch_idx, (inputs, targets) in tqdm(enumerate(train_loader),
+                                             total=len(dataset) // config["batch_size"]):
+        inputs, targets = inputs.to(device), targets.to(device)
+        optimizer.zero_grad()
+        with torch.cuda.amp.autocast(enabled=True, dtype=torch.bfloat16):
+            outputs = model(inputs)
+            loss = criterion(outputs, targets)
+        loss.backward()
+        optimizer.step()
+    if scheduler is not None:
+        scheduler.step()
+
+# test
+@torch.no_grad()
+def test(model=model):
+    model.eval()
+    all_targets = []
+    all_predicts = []
+    test_loss = 0
+    correct = 0
+    total = 0
+    for batch_idx, (inputs, targets) in tqdm(enumerate(test_loader),
+                                             total=len(test_loader.dataset) // config["batch_size"]):
+        inputs, targets = inputs.to(device), targets.to(device)
+        with torch.cuda.amp.autocast(enabled=False, dtype=torch.bfloat16):
+            outputs = model(inputs)
+            loss = criterion(outputs, targets)
+        # to logging losses
+        all_targets.extend(targets.flatten().tolist())
+        test_loss += loss.item()
+        _, predicts = outputs.max(1)
+        all_predicts.extend(predicts.flatten().tolist())
+        total += targets.size(0)
+        correct += predicts.eq(targets).sum().item()
+    loss = test_loss / (batch_idx + 1)
+    acc = correct / total
+    print(f"Loss: {loss:.4f} | Acc: {acc:.4f}\n")
+    model.train()
+    return loss, acc, all_targets, all_predicts
+
+# save train
+def save_train(model=model, optimizer=optimizer):
+    model.train()
+    for batch_idx, (inputs, targets) in enumerate(train_loader):
+        inputs, targets = inputs.to(device), targets.to(device)
+        optimizer.zero_grad()
+        with torch.cuda.amp.autocast(enabled=True, dtype=torch.bfloat16):
+            outputs = model(inputs)
+            loss = criterion(outputs, targets)
+        loss.backward()
+        optimizer.step()
+        # Save checkpoint
+        if batch_idx % (len(dataset) // train_loader.batch_size // config["total_save_number"]) == 0:
+            _, acc, _, _ = test(model=model)
+            if not os.path.isdir('checkpoint'):
+                os.mkdir('checkpoint')
+            save_state = {key: value.cpu().to(torch.float32) for key, value in model.state_dict().items()}
+            torch.save(save_state, f"checkpoint/{str(batch_idx).zfill(4)}_acc{acc:.4f}_seed{seed:04d}_{config['tag']}.pth")
+            print("save:", f"checkpoint/{str(batch_idx).zfill(4)}_acc{acc:.4f}_seed{seed:04d}_{config['tag']}.pth")
+
+
+
+
+# main
+if __name__ == '__main__':
+    test(model=model)
+    for epoch in range(config["epochs"]):
+        train(model=model, optimizer=optimizer, scheduler=scheduler)
+        test(model=model)
+    save_train(model=model, optimizer=optimizer)

dataset/cifar10_mobilenetv3/model.py

@@ -0,0 +1,21 @@
+import torch.nn as nn
+import timm
+
+
+def Model():
+    model = timm.create_model("mobilenetv3_large_100", pretrained=True)
+    model.classifier = nn.Linear(1280, 10)
+    for name, param in model.named_parameters():
+        if "bn" in name:
+            # print(f"freeze {name}")
+            param.requires_grad = False
+    return model, model.classifier
+
+
+if __name__ == "__main__":
+    model, _ = Model()
+    print(model)
+    num_param = 0
+    for v in model.parameters():
+        num_param += v.numel()
+    print("num_param:", num_param)

dataset/cifar10_mobilenetv3/test.py

@@ -0,0 +1,28 @@
+import os
+import sys
+if __name__ == "__main__":
+    from train import *
+else:  # relative import
+    from .train import *
+
+
+
+
+try:
+    test_item = sys.argv[1]
+except IndexError:
+    assert __name__ == "__main__"
+    test_item = "./checkpoint"
+test_items = []
+if os.path.isdir(test_item):
+    for item in os.listdir(test_item):
+        item = os.path.join(test_item, item)
+        test_items.append(item)
+elif os.path.isfile(test_item):
+    test_items.append(test_item)
+
+
+for item in test_items:
+    state = torch.load(item, map_location="cpu")
+    model.load_state_dict({key: value.to(torch.float32).to(device) for key, value in state.items()})
+    loss, acc, all_targets, all_predicts = test(model=model)

dataset/cifar10_mobilenetv3/train.py

@@ -0,0 +1,199 @@
+# set global seed
+import random
+import numpy as np
+import torch
+seed = SEED = 20
+torch.manual_seed(seed)
+torch.cuda.manual_seed(seed)
+torch.cuda.manual_seed_all(seed)
+torch.backends.cudnn.deterministic = True
+torch.backends.cudnn.benchmark = True
+np.random.seed(seed)
+random.seed(seed)
+
+
+try:  # relative import
+    from model import Model
+except ImportError:
+    from .model import Model
+
+# import
+import torch.nn as nn
+from torch import optim
+from torch.optim import lr_scheduler
+from torch.utils.data import DataLoader
+import torchvision.transforms as transforms
+from torchvision.datasets import CIFAR10 as Dataset
+from tqdm.auto import tqdm
+import os
+import warnings
+warnings.filterwarnings("ignore", category=UserWarning)
+
+# load additional config
+import json
+config_file = os.path.join(os.path.dirname(os.path.dirname(__file__)), "config.json")
+with open(config_file, "r") as f:
+    additional_config = json.load(f)
+
+
+
+
+# config
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+config = {
+    "dataset_root": "from_additional_config",
+    "batch_size": 500 if __name__ == "__main__" else 200,
+    "num_workers": 4,
+    "learning_rate": 3e-3,
+    "weight_decay": 0.1,
+    "epochs": 5,
+    "save_learning_rate": 1e-6,
+    "total_save_number": 50,
+    "tag": os.path.basename(os.path.dirname(__file__)),
+}
+config.update(additional_config)
+
+
+
+
+# Data
+dataset = Dataset(
+    root=config["dataset_root"],
+    download=True,
+    train=True,
+    transform=transforms.Compose([
+        transforms.Resize(224),
+        transforms.RandomCrop(224),
+        transforms.RandomHorizontalFlip(),
+        transforms.AutoAugment(policy=transforms.AutoAugmentPolicy("cifar10")),
+        transforms.ToTensor(),
+        transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
+    ])
+)
+train_loader = DataLoader(
+    dataset=dataset,
+    batch_size=config["batch_size"],
+    num_workers=config["num_workers"],
+    shuffle=True,
+    drop_last=True,
+    pin_memory=True,
+    persistent_workers=True,
+)
+test_loader = DataLoader(
+    dataset=Dataset(
+        root=config["dataset_root"],
+        download=True,
+        train=False,
+        transform=transforms.Compose([
+            transforms.Resize(224),
+            transforms.CenterCrop(224),
+            transforms.ToTensor(),
+            transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
+        ])),
+    batch_size=config["batch_size"],
+    num_workers=config["num_workers"],
+    shuffle=False,
+    pin_memory=True,
+    persistent_workers=True,
+    pin_memory_device="cuda",
+)
+
+# Model
+model, head = Model()
+model = model.to(device)
+criterion = nn.CrossEntropyLoss()
+pre_optimizer = optim.AdamW(
+    head.parameters(),
+    lr=0.05,
+    weight_decay=0.01,
+)
+optimizer = optim.AdamW(
+    model.parameters(),
+    lr=config["learning_rate"],
+    weight_decay=config["weight_decay"],
+)
+scheduler = lr_scheduler.CosineAnnealingLR(
+    optimizer,
+    T_max=config["epochs"],
+    eta_min=config["save_learning_rate"],
+)
+
+
+
+
+# Training
+def train(model=model, optimizer=optimizer, scheduler=scheduler):
+    model.train()
+    for batch_idx, (inputs, targets) in tqdm(enumerate(train_loader),
+                                             total=len(dataset) // config["batch_size"]):
+        inputs, targets = inputs.to(device), targets.to(device)
+        optimizer.zero_grad()
+        with torch.cuda.amp.autocast(enabled=True, dtype=torch.bfloat16):
+            outputs = model(inputs)
+            loss = criterion(outputs, targets)
+        loss.backward()
+        optimizer.step()
+    if scheduler is not None:
+        scheduler.step()
+
+# test
+@torch.no_grad()
+def test(model=model):
+    model.eval()
+    all_targets = []
+    all_predicts = []
+    test_loss = 0
+    correct = 0
+    total = 0
+    for batch_idx, (inputs, targets) in tqdm(enumerate(test_loader),
+                                             total=len(test_loader.dataset) // config["batch_size"]):
+        inputs, targets = inputs.to(device), targets.to(device)
+        with torch.cuda.amp.autocast(enabled=False, dtype=torch.bfloat16):
+            outputs = model(inputs)
+            loss = criterion(outputs, targets)
+        # to logging losses
+        all_targets.extend(targets.flatten().tolist())
+        test_loss += loss.item()
+        _, predicts = outputs.max(1)
+        all_predicts.extend(predicts.flatten().tolist())
+        total += targets.size(0)
+        correct += predicts.eq(targets).sum().item()
+    loss = test_loss / (batch_idx + 1)
+    acc = correct / total
+    print(f"Loss: {loss:.4f} | Acc: {acc:.4f}\n")
+    model.train()
+    return loss, acc, all_targets, all_predicts
+
+# save train
+def save_train(model=model, optimizer=optimizer):
+    model.train()
+    for batch_idx, (inputs, targets) in enumerate(train_loader):
+        inputs, targets = inputs.to(device), targets.to(device)
+        optimizer.zero_grad()
+        with torch.cuda.amp.autocast(enabled=True, dtype=torch.bfloat16):
+            outputs = model(inputs)
+            loss = criterion(outputs, targets)
+        loss.backward()
+        optimizer.step()
+        # Save checkpoint
+        if batch_idx % (len(dataset) // train_loader.batch_size // config["total_save_number"]) == 0:
+            _, acc, _, _ = test(model=model)
+            if not os.path.isdir('checkpoint'):
+                os.mkdir('checkpoint')
+            save_state = {key: value.cpu().to(torch.float32) for key, value in model.state_dict().items()}
+            torch.save(save_state, f"checkpoint/{str(batch_idx).zfill(4)}_acc{acc:.4f}_seed{seed:04d}_{config['tag']}.pth")
+            print("save:", f"checkpoint/{str(batch_idx).zfill(4)}_acc{acc:.4f}_seed{seed:04d}_{config['tag']}.pth")
+
+
+
+
+# main
+if __name__ == '__main__':
+    test(model=model)
+    for _ in range(1):
+        train(model=model, optimizer=pre_optimizer)
+        test(model=model)
+    for epoch in range(config["epochs"]):
+        train(model=model, optimizer=optimizer, scheduler=scheduler)
+        test(model=model)
+    save_train(model=model, optimizer=optimizer)

dataset/cifar10_resnet18/model.py

@@ -0,0 +1,17 @@
+import torch.nn as nn
+import timm
+
+
+def Model():
+    model = timm.create_model("resnet18", pretrained=True)
+    model.fc = nn.Linear(512, 10)
+    return model, model.fc
+
+
+if __name__ == "__main__":
+    model, _ = Model()
+    print(model)
+    num_param = 0
+    for v in model.parameters():
+        num_param += v.numel()
+    print("num_param:", num_param)

dataset/cifar10_resnet18/test.py

@@ -0,0 +1,28 @@
+import os
+import sys
+if __name__ == "__main__":
+    from train import *
+else:  # relative import
+    from .train import *
+
+
+
+
+try:
+    test_item = sys.argv[1]
+except IndexError:
+    assert __name__ == "__main__"
+    test_item = "./checkpoint"
+test_items = []
+if os.path.isdir(test_item):
+    for item in os.listdir(test_item):
+        item = os.path.join(test_item, item)
+        test_items.append(item)
+elif os.path.isfile(test_item):
+    test_items.append(test_item)
+
+
+for item in test_items:
+    state = torch.load(item, map_location="cpu")
+    model.load_state_dict({key: value.to(torch.float32).to(device) for key, value in state.items()})
+    loss, acc, all_targets, all_predicts = test(model=model)

dataset/cifar10_resnet18/train.py

@@ -0,0 +1,191 @@
+# set global seed
+import random
+import numpy as np
+import torch
+seed = SEED = 20
+torch.manual_seed(seed)
+torch.cuda.manual_seed(seed)
+torch.cuda.manual_seed_all(seed)
+torch.backends.cudnn.deterministic = True
+torch.backends.cudnn.benchmark = True
+np.random.seed(seed)
+random.seed(seed)
+
+
+try:  # relative import
+    from model import Model
+except ImportError:
+    from .model import Model
+
+# import
+import torch.nn as nn
+from torch import optim
+from torch.optim import lr_scheduler
+from torch.utils.data import DataLoader
+import torchvision.transforms as transforms
+from torchvision.datasets import CIFAR10 as Dataset
+from tqdm.auto import tqdm
+import os
+import warnings
+warnings.filterwarnings("ignore", category=UserWarning)
+
+# load additional config
+import json
+config_file = os.path.join(os.path.dirname(os.path.dirname(__file__)), "config.json")
+with open(config_file, "r") as f:
+    additional_config = json.load(f)
+
+
+
+
+# config
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+config = {
+    "dataset_root": "from_additional_config",
+    "batch_size": 500 if __name__ == "__main__" else 200,
+    "num_workers": 32,
+    "learning_rate": 3e-3,
+    "weight_decay": 0.1,
+    "epochs": 50,
+    "save_learning_rate": 1e-5,
+    "total_save_number": 50,
+    "tag": os.path.basename(os.path.dirname(__file__)),
+}
+config.update(additional_config)
+
+
+
+
+# Data
+dataset = Dataset(
+    root=config["dataset_root"],
+    download=True,
+    train=True,
+    transform=transforms.Compose([
+        transforms.Resize(64),
+        transforms.RandomCrop(64),
+        transforms.RandomHorizontalFlip(),
+        transforms.AutoAugment(policy=transforms.AutoAugmentPolicy("cifar10")),
+        transforms.ToTensor(),
+        transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
+    ])
+)
+train_loader = DataLoader(
+    dataset=dataset,
+    batch_size=config["batch_size"],
+    num_workers=config["num_workers"],
+    shuffle=True,
+    drop_last=True,
+    pin_memory=True,
+    persistent_workers=True,
+)
+test_loader = DataLoader(
+    dataset=Dataset(
+        root=config["dataset_root"],
+        download=True,
+        train=False,
+        transform=transforms.Compose([
+            transforms.Resize(64),
+            transforms.CenterCrop(64),
+            transforms.ToTensor(),
+            transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
+        ])),
+    batch_size=config["batch_size"],
+    num_workers=config["num_workers"],
+    shuffle=False,
+    pin_memory=True,
+    persistent_workers=True,
+    pin_memory_device="cuda",
+)
+
+# Model
+model, head = Model()
+model = model.to(device)
+criterion = nn.CrossEntropyLoss()
+optimizer = optim.AdamW(
+    model.parameters(),
+    lr=config["learning_rate"],
+    weight_decay=config["weight_decay"],
+)
+scheduler = lr_scheduler.CosineAnnealingLR(
+    optimizer,
+    T_max=config["epochs"],
+    eta_min=config["save_learning_rate"],
+)
+
+
+
+
+# Training
+def train(model=model, optimizer=optimizer, scheduler=scheduler):
+    model.train()
+    for batch_idx, (inputs, targets) in tqdm(enumerate(train_loader),
+                                             total=len(dataset) // config["batch_size"]):
+        inputs, targets = inputs.to(device), targets.to(device)
+        optimizer.zero_grad()
+        with torch.cuda.amp.autocast(enabled=True, dtype=torch.bfloat16):
+            outputs = model(inputs)
+            loss = criterion(outputs, targets)
+        loss.backward()
+        optimizer.step()
+    if scheduler is not None:
+        scheduler.step()
+
+# test
+@torch.no_grad()
+def test(model=model):
+    model.eval()
+    all_targets = []
+    all_predicts = []
+    test_loss = 0
+    correct = 0
+    total = 0
+    for batch_idx, (inputs, targets) in tqdm(enumerate(test_loader),
+                                             total=len(test_loader.dataset) // config["batch_size"]):
+        inputs, targets = inputs.to(device), targets.to(device)
+        with torch.cuda.amp.autocast(enabled=False, dtype=torch.bfloat16):
+            outputs = model(inputs)
+            loss = criterion(outputs, targets)
+        # to logging losses
+        all_targets.extend(targets.flatten().tolist())
+        test_loss += loss.item()
+        _, predicts = outputs.max(1)
+        all_predicts.extend(predicts.flatten().tolist())
+        total += targets.size(0)
+        correct += predicts.eq(targets).sum().item()
+    loss = test_loss / (batch_idx + 1)
+    acc = correct / total
+    print(f"Loss: {loss:.4f} | Acc: {acc:.4f}\n")
+    model.train()
+    return loss, acc, all_targets, all_predicts
+
+# save train
+def save_train(model=model, optimizer=optimizer):
+    model.train()
+    for batch_idx, (inputs, targets) in enumerate(train_loader):
+        inputs, targets = inputs.to(device), targets.to(device)
+        optimizer.zero_grad()
+        with torch.cuda.amp.autocast(enabled=True, dtype=torch.bfloat16):
+            outputs = model(inputs)
+            loss = criterion(outputs, targets)
+        loss.backward()
+        optimizer.step()
+        # Save checkpoint
+        if batch_idx % (len(dataset) // train_loader.batch_size // config["total_save_number"]) == 0:
+            _, acc, _, _ = test(model=model)
+            if not os.path.isdir('checkpoint'):
+                os.mkdir('checkpoint')
+            save_state = {key: value.cpu().to(torch.float32) for key, value in model.state_dict().items()}
+            torch.save(save_state, f"checkpoint/{str(batch_idx).zfill(4)}_acc{acc:.4f}_seed{seed:04d}_{config['tag']}.pth")
+            print("save:", f"checkpoint/{str(batch_idx).zfill(4)}_acc{acc:.4f}_seed{seed:04d}_{config['tag']}.pth")
+
+
+
+
+# main
+if __name__ == '__main__':
+    test(model=model)
+    for epoch in range(config["epochs"]):
+        train(model=model, optimizer=optimizer, scheduler=scheduler)
+        test(model=model)
+    save_train(model=model, optimizer=optimizer)

dataset/cifar10_vitbase/model.py

@@ -0,0 +1,17 @@
+import torch.nn as nn
+import timm
+
+
+def Model():
+    model = timm.create_model("vit_base_patch16_224", pretrained=True)
+    model.head = nn.Linear(768, 10)
+    return model, model.head
+
+
+if __name__ == "__main__":
+    model, _ = Model()
+    print(model)
+    num_param = 0
+    for v in model.parameters():
+        num_param += v.numel()
+    print("num_param:", num_param)

dataset/cifar10_vitbase/test.py

@@ -0,0 +1,28 @@
+import os
+import sys
+if __name__ == "__main__":
+    from train import *
+else:  # relative import
+    from .train import *
+
+
+
+
+try:
+    test_item = sys.argv[1]
+except IndexError:
+    assert __name__ == "__main__"
+    test_item = "./checkpoint"
+test_items = []
+if os.path.isdir(test_item):
+    for item in os.listdir(test_item):
+        item = os.path.join(test_item, item)
+        test_items.append(item)
+elif os.path.isfile(test_item):
+    test_items.append(test_item)
+
+
+for item in test_items:
+    state = torch.load(item, map_location="cpu")
+    model.load_state_dict({key: value.to(torch.float32).to(device) for key, value in state.items()})
+    loss, acc, all_targets, all_predicts = test(model=model)

dataset/cifar10_vitbase/train.py

@@ -0,0 +1,199 @@
+# set global seed
+import random
+import numpy as np
+import torch
+seed = SEED = 20
+torch.manual_seed(seed)
+torch.cuda.manual_seed(seed)
+torch.cuda.manual_seed_all(seed)
+torch.backends.cudnn.deterministic = True
+torch.backends.cudnn.benchmark = True
+np.random.seed(seed)
+random.seed(seed)
+
+
+try:  # relative import
+    from model import Model
+except ImportError:
+    from .model import Model
+
+# import
+import torch.nn as nn
+from torch import optim
+from torch.optim import lr_scheduler
+from torch.utils.data import DataLoader
+import torchvision.transforms as transforms
+from torchvision.datasets import CIFAR10 as Dataset
+from tqdm.auto import tqdm
+import os
+import warnings
+warnings.filterwarnings("ignore", category=UserWarning)
+
+# load additional config
+import json
+config_file = os.path.join(os.path.dirname(os.path.dirname(__file__)), "config.json")
+with open(config_file, "r") as f:
+    additional_config = json.load(f)
+
+
+
+
+# config
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+config = {
+    "dataset_root": "from_additional_config",
+    "batch_size": 500 if __name__ == "__main__" else 200,
+    "num_workers": 32,
+    "learning_rate": 3e-5,
+    "weight_decay": 0.1,
+    "epochs": 7,
+    "save_learning_rate": 1e-5,
+    "total_save_number": 50,
+    "tag": os.path.basename(os.path.dirname(__file__)),
+}
+config.update(additional_config)
+
+
+
+
+# Data
+dataset = Dataset(
+    root=config["dataset_root"],
+    download=True,
+    train=True,
+    transform=transforms.Compose([
+        transforms.Resize(224),
+        transforms.RandomCrop(224),
+        transforms.RandomHorizontalFlip(),
+        transforms.AutoAugment(policy=transforms.AutoAugmentPolicy("cifar10")),
+        transforms.ToTensor(),
+        transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
+    ])
+)
+train_loader = DataLoader(
+    dataset=dataset,
+    batch_size=config["batch_size"],
+    num_workers=config["num_workers"],
+    shuffle=True,
+    drop_last=True,
+    pin_memory=True,
+    persistent_workers=True,
+)
+test_loader = DataLoader(
+    dataset=Dataset(
+        root=config["dataset_root"],
+        download=True,
+        train=False,
+        transform=transforms.Compose([
+            transforms.Resize(224),
+            transforms.CenterCrop(224),
+            transforms.ToTensor(),
+            transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
+        ])),
+    batch_size=config["batch_size"],
+    num_workers=config["num_workers"],
+    shuffle=False,
+    pin_memory=True,
+    persistent_workers=True,
+    pin_memory_device="cuda",
+)
+
+# Model
+model, head = Model()
+model = model.to(device)
+criterion = nn.CrossEntropyLoss()
+pre_optimizer = optim.AdamW(
+    head.parameters(),
+    lr=0.05,
+    weight_decay=0.01,
+)
+optimizer = optim.AdamW(
+    model.parameters(),
+    lr=config["learning_rate"],
+    weight_decay=config["weight_decay"],
+)
+scheduler = lr_scheduler.CosineAnnealingLR(
+    optimizer,
+    T_max=config["epochs"],
+    eta_min=config["save_learning_rate"],
+)
+
+
+
+
+# Training
+def train(model=model, optimizer=optimizer, scheduler=scheduler):
+    model.train()
+    for batch_idx, (inputs, targets) in tqdm(enumerate(train_loader),
+                                             total=len(dataset) // config["batch_size"]):
+        inputs, targets = inputs.to(device), targets.to(device)
+        optimizer.zero_grad()
+        with torch.cuda.amp.autocast(enabled=True, dtype=torch.bfloat16):
+            outputs = model(inputs)
+            loss = criterion(outputs, targets)
+        loss.backward()
+        optimizer.step()
+    if scheduler is not None:
+        scheduler.step()
+
+# test
+@torch.no_grad()
+def test(model=model):
+    model.eval()
+    all_targets = []
+    all_predicts = []
+    test_loss = 0
+    correct = 0
+    total = 0
+    for batch_idx, (inputs, targets) in tqdm(enumerate(test_loader),
+                                             total=len(test_loader.dataset) // config["batch_size"]):
+        inputs, targets = inputs.to(device), targets.to(device)
+        with torch.cuda.amp.autocast(enabled=False, dtype=torch.bfloat16):
+            outputs = model(inputs)
+            loss = criterion(outputs, targets)
+        # to logging losses
+        all_targets.extend(targets.flatten().tolist())
+        test_loss += loss.item()
+        _, predicts = outputs.max(1)
+        all_predicts.extend(predicts.flatten().tolist())
+        total += targets.size(0)
+        correct += predicts.eq(targets).sum().item()
+    loss = test_loss / (batch_idx + 1)
+    acc = correct / total
+    print(f"Loss: {loss:.4f} | Acc: {acc:.4f}\n")
+    model.train()
+    return loss, acc, all_targets, all_predicts
+
+# save train
+def save_train(model=model, optimizer=optimizer):
+    model.train()
+    for batch_idx, (inputs, targets) in enumerate(train_loader):
+        inputs, targets = inputs.to(device), targets.to(device)
+        optimizer.zero_grad()
+        with torch.cuda.amp.autocast(enabled=True, dtype=torch.bfloat16):
+            outputs = model(inputs)
+            loss = criterion(outputs, targets)
+        loss.backward()
+        optimizer.step()
+        # Save checkpoint
+        if batch_idx % (len(dataset) // train_loader.batch_size // config["total_save_number"]) == 0:
+            _, acc, _, _ = test(model=model)
+            if not os.path.isdir('checkpoint'):
+                os.mkdir('checkpoint')
+            save_state = {key: value.cpu().to(torch.float32) for key, value in model.state_dict().items()}
+            torch.save(save_state, f"checkpoint/{str(batch_idx).zfill(4)}_acc{acc:.4f}_seed{seed:04d}_{config['tag']}.pth")
+            print("save:", f"checkpoint/{str(batch_idx).zfill(4)}_acc{acc:.4f}_seed{seed:04d}_{config['tag']}.pth")
+
+
+
+
+# main
+if __name__ == '__main__':
+    test(model=model)
+    for _ in range(3):
+        train(model=model, optimizer=pre_optimizer)
+        test(model=model)
+    for epoch in range(config["epochs"]):
+        train(model=model, optimizer=optimizer, scheduler=scheduler)
+        test(model=model)
+    save_train(model=model, optimizer=optimizer)

dataset/condition_classinput_inference/dataset.py

@@ -0,0 +1,41 @@
+import re
+import sys
+from torch.utils.data import Dataset
+from torchvision.datasets import CIFAR10
+import torchvision.transforms as transforms
+
+
+class BinaryClassifierDataset(Dataset):
+    def __init__(self, root, train, optimize_class: list):
+        self.optimize_class = optimize_class
+        self.dataset = CIFAR10(
+            root=root,
+            train=train,
+            download=True,
+            transform=transforms.Compose([
+                transforms.Resize(224),
+                transforms.RandomHorizontalFlip(),
+                transforms.AutoAugment(policy=transforms.AutoAugmentPolicy("cifar10")),
+                transforms.ToTensor(),
+                transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
+            ])
+        )
+
+    def __getitem__(self, index):
+        img, origin_target = self.dataset[index]
+        target = 1 if origin_target in self.optimize_class else 0
+        return img, target
+
+    def __len__(self):
+        return self.dataset.__len__()
+
+
+def get_optimize_class():
+    try:  # get string
+        string = sys.argv[1]
+    except IndexError:
+        RuntimeError("sys.argv[1] not found")
+    class_int_string = str(re.search(r'class(\d+)', string).group(1)).zfill(4)
+    one_hot_string = bin(int(class_int_string))[2:].zfill(10)
+    optimize_class = [index for index, i in enumerate(one_hot_string) if i == "1"]
+    return list(optimize_class), class_int_string

dataset/condition_classinput_inference/model.py

@@ -0,0 +1,25 @@
+import torch
+import torch.nn as nn
+import timm
+
+
+def Model():
+    model = timm.create_model("vit_tiny_patch16_224", pretrained=True)
+    model.head = nn.Sequential(
+        nn.Linear(192, 192, bias=True),
+        nn.SiLU(),
+        nn.Linear(192, 2, bias=False),
+    )
+    for param in model.head.parameters():
+        param = nn.Parameter(torch.ones_like(param) / 192)
+        param.requires_grad = True
+    return model, model.head
+
+
+if __name__ == "__main__":
+    model, _ = Model()
+    print(model)
+    num_param = 0
+    for v in model.parameters():
+        num_param += v.numel()
+    print("num_param:", num_param)

dataset/condition_classinput_inference/test.py

@@ -0,0 +1,30 @@
+import os
+import sys
+if __name__ == "__main__":
+    from train import *
+else:  # relative import
+    from .train import *
+
+
+
+
+try:
+    test_item = sys.argv[1]
+except IndexError:
+    assert __name__ == "__main__"
+    test_item = "./checkpoint_test"
+test_items = []
+if os.path.isdir(test_item):
+    for item in os.listdir(test_item):
+        item = os.path.join(test_item, item)
+        test_items.append(item)
+elif os.path.isfile(test_item):
+    test_items.append(test_item)
+
+
+
+
+for item in test_items:
+    state = torch.load(item, map_location="cpu")
+    model.load_state_dict({key: value.to(torch.float32).to(device) for key, value in state.items()})
+    loss, acc, all_targets, all_predicts = test(model=model)

dataset/condition_classinput_inference/train.py

@@ -0,0 +1,209 @@
+# set global seed
+import random
+import numpy as np
+import torch
+seed = SEED = 20
+torch.manual_seed(seed)
+torch.cuda.manual_seed(seed)
+torch.cuda.manual_seed_all(seed)
+torch.backends.cudnn.deterministic = True
+torch.backends.cudnn.benchmark = True
+np.random.seed(seed)
+random.seed(seed)
+
+
+try:  # relative import
+    from model import Model
+    from dataset import BinaryClassifierDataset as Dataset
+    from dataset import get_optimize_class
+except ImportError:
+    from .model import Model
+    from .dataset import BinaryClassifierDataset as Dataset
+    from .dataset import get_optimize_class
+
+# import
+import torch.nn as nn
+from torch import optim
+from torch.optim import lr_scheduler
+from torch.utils.data import DataLoader
+from torch.nn import functional as F
+import os
+import sys
+import warnings
+warnings.filterwarnings("ignore", category=UserWarning)
+
+# load additional config
+import json
+config_file = os.path.join(os.path.dirname(os.path.dirname(__file__)), "config.json")
+with open(config_file, "r") as f:
+    additional_config = json.load(f)
+
+
+
+
+# config
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+config = {
+    "dataset_root": "from_additional_config",
+    "batch_size": 500 if __name__ == "__main__" else 50,
+    "num_workers": 16,
+    "pre_learning_rate": 0.01,
+    "learning_rate": 1e-4,
+    "pre_epochs": 2,
+    "epochs": 13,
+    "weight_decay": 0.1,
+    "save_learning_rate": 2e-5,
+    "total_save_number": 5,
+    "tag": os.path.basename(os.path.dirname(__file__)),
+    "optimize_class": get_optimize_class()[0],
+    "optimize_class_int": get_optimize_class()[1],
+}
+config.update(additional_config)
+print("Training/Testing:", config["optimize_class"])
+
+
+
+
+# Data
+dataset = Dataset(
+    root=config["dataset_root"],
+    train=True,
+    optimize_class=config["optimize_class"],
+)
+train_loader = DataLoader(
+    dataset=dataset,
+    batch_size=config["batch_size"],
+    num_workers=config["num_workers"],
+    shuffle=True,
+    drop_last=True,
+    pin_memory=True,
+    persistent_workers=True,
+)
+test_loader = DataLoader(
+    dataset=Dataset(
+        root=config["dataset_root"],
+        train=False,
+        optimize_class=config["optimize_class"],
+    ),
+    batch_size=config["batch_size"],
+    num_workers=config["num_workers"],
+    shuffle=False,
+)
+
+# Model
+model, head = Model()
+model = model.to(device)
+class FocalLoss(nn.Module):
+    def __init__(self, weight=None, gamma=2):
+        super(FocalLoss, self).__init__()
+        self.weight = weight
+        self.gamma = gamma
+    def forward(self, input, target):
+        ce_loss = F.cross_entropy(input, target, reduction='none', weight=self.weight)
+        pt = torch.exp(-ce_loss)
+        focal_loss = (1 - pt) ** self.gamma * ce_loss
+        return focal_loss.mean()
+criterion = FocalLoss()
+
+# Optimizer
+head_optimizer = optim.AdamW(
+    head.parameters(),
+    lr=config["pre_learning_rate"],
+    weight_decay=config["weight_decay"],
+)
+optimizer = optim.AdamW(
+    model.parameters(),
+    lr=config["learning_rate"],
+    weight_decay=config["weight_decay"],
+)
+scheduler = lr_scheduler.CosineAnnealingLR(
+    optimizer,
+    T_max=config["epochs"],
+    eta_min=config["save_learning_rate"],
+)
+
+
+
+
+# Training
+def train(model=model, optimizer=optimizer, scheduler=scheduler):
+    model.train()
+    for batch_idx, (inputs, targets) in enumerate(train_loader):
+        inputs, targets = inputs.to(device), targets.to(device)
+        optimizer.zero_grad()
+        with torch.cuda.amp.autocast(enabled=False, dtype=torch.bfloat16):
+            outputs = model(inputs)
+            loss = criterion(outputs, targets)
+        loss.backward()
+        optimizer.step()
+    if scheduler is not None:
+        scheduler.step()
+
+# test
+@torch.no_grad()
+def test(model=model):
+    model.eval()
+    all_targets = []
+    all_predicts = []
+    test_loss = 0
+    correct = 0
+    total = 0
+    for batch_idx, (inputs, targets) in enumerate(test_loader):
+        inputs, targets = inputs.to(device), targets.to(device)
+        with torch.cuda.amp.autocast(enabled=False, dtype=torch.bfloat16):
+            outputs = model(inputs)
+            loss = criterion(outputs, targets)
+        # to logging losses
+        all_targets.extend(targets.flatten().tolist())
+        test_loss += loss.item()
+        _, predicts = outputs.max(1)
+        all_predicts.extend(predicts.flatten().tolist())
+        total += targets.size(0)
+        correct += predicts.eq(targets).sum().item()
+    loss = test_loss / (batch_idx + 1)
+    acc = correct / total
+    print(f"Loss: {loss:.4f} | Acc: {acc:.4f}\n")
+    model.train()
+    return loss, acc, all_targets, all_predicts
+
+# save train
+def save_train(model=model, optimizer=optimizer):
+    data_loader = DataLoader(
+        dataset=dataset,
+        batch_size=min(len(dataset) // config["total_save_number"], config["batch_size"]),
+        num_workers=config["num_workers"],
+        shuffle=True,
+        drop_last=True,
+    )
+    model.train()
+    for batch_idx, (inputs, targets) in enumerate(data_loader):
+        inputs, targets = inputs.to(device), targets.to(device)
+        optimizer.zero_grad()
+        with torch.cuda.amp.autocast(enabled=False, dtype=torch.bfloat16):
+            outputs = model(inputs)
+            loss = criterion(outputs, targets)
+        loss.backward()
+        optimizer.step()
+        # Save checkpoint
+        _, acc, _, _ = test(model=model)
+        if not os.path.isdir('checkpoint'):
+            os.mkdir('checkpoint')
+        save_state = {key: value.cpu().to(torch.float32) for key, value in model.state_dict().items()}
+        torch.save(save_state, f"checkpoint/{str(batch_idx).zfill(4)}_acc{acc:.4f}_class{config['optimize_class_int']}_{config['tag']}.pth")
+        print("save:", f"checkpoint/{str(batch_idx).zfill(4)}_acc{acc:.4f}_class{config['optimize_class_int']}_{config['tag']}.pth")
+        # exit loop
+        if batch_idx+1 == config["total_save_number"]:
+            break
+
+
+
+
+# main
+if __name__ == '__main__':
+    for epoch in range(config["pre_epochs"]):
+        train(model=model, optimizer=head_optimizer, scheduler=None)
+        # test(model=model)
+    for epoch in range(config["epochs"]):
+        train(model=model, optimizer=optimizer, scheduler=scheduler)
+        # test(model=model)
+    save_train(model=model, optimizer=optimizer)

dataset/condition_classinput_vittiny/dataset.py

@@ -0,0 +1,41 @@
+import re
+import sys
+from torch.utils.data import Dataset
+from torchvision.datasets import CIFAR10
+import torchvision.transforms as transforms
+
+
+class BinaryClassifierDataset(Dataset):
+    def __init__(self, root, train, optimize_class: list):
+        self.optimize_class = optimize_class
+        self.dataset = CIFAR10(
+            root=root,
+            train=train,
+            download=True,
+            transform=transforms.Compose([
+                transforms.Resize(224),
+                transforms.RandomHorizontalFlip(),
+                transforms.AutoAugment(policy=transforms.AutoAugmentPolicy("cifar10")),
+                transforms.ToTensor(),
+                transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
+            ])
+        )
+
+    def __getitem__(self, index):
+        img, origin_target = self.dataset[index]
+        target = 1 if origin_target in self.optimize_class else 0
+        return img, target
+
+    def __len__(self):
+        return self.dataset.__len__()
+
+
+def get_optimize_class():
+    try:  # get string
+        string = sys.argv[1]
+    except IndexError:
+        RuntimeError("sys.argv[1] not found")
+    class_int_string = str(re.search(r'class(\d+)', string).group(1)).zfill(4)
+    one_hot_string = bin(int(class_int_string))[2:].zfill(10)
+    optimize_class = [index for index, i in enumerate(one_hot_string) if i == "1"]
+    return list(optimize_class), class_int_string

dataset/condition_classinput_vittiny/detail.py

@@ -0,0 +1,58 @@
+import os
+import sys
+if __name__ == "__main__":
+    from train import *
+else:  # relative import
+    from .train import *
+from torchvision.datasets import CIFAR10
+from torchvision import transforms
+
+
+
+
+try:
+    test_item = sys.argv[1]
+except IndexError:
+    assert __name__ == "__main__"
+    test_item = "./generated"
+test_items = []
+if os.path.isdir(test_item):
+    for item in os.listdir(test_item):
+        item = os.path.join(test_item, item)
+        test_items.append(item)
+elif os.path.isfile(test_item):
+    test_items.append(test_item)
+
+
+
+
+original_dataset = CIFAR10(
+    root=config["dataset_root"],
+    train=False,
+    download=True,
+    transform=transforms.Compose([
+        transforms.Resize(224),
+        transforms.ToTensor(),
+        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
+    ])
+)
+original_targets = [original_dataset[i][1] for i in range(len(original_dataset))]
+original_targets = torch.tensor(original_targets, dtype=torch.long)
+
+
+
+
+for item in test_items:
+    state = torch.load(item, map_location="cpu")
+    model.load_state_dict({key: value.to(torch.float32).to(device) for key, value in state.items()})
+    loss, acc, all_targets, all_predicts = test(model=model)
+    all_targets, all_predicts = torch.tensor(all_targets), torch.tensor(all_predicts)
+
+    for class_idx in range(10):
+        class_mask = torch.where(original_targets == class_idx, 1, 0)
+        total_number = torch.sum(class_mask)
+        correct = torch.where(all_targets == all_predicts, 1, 0)
+        class_correct = class_mask * correct
+        correct_number = torch.sum(class_correct)
+        class_acc = correct_number.item() / total_number.item()
+        print(f"class{class_idx}:", class_acc)

dataset/condition_classinput_vittiny/finetune.py

@@ -0,0 +1,215 @@
+# set global seed
+import time
+print("time stamp:", time.time())
+import random
+import numpy as np
+import torch
+seed = SEED = 20
+torch.manual_seed(seed)
+torch.cuda.manual_seed(seed)
+torch.cuda.manual_seed_all(seed)
+torch.backends.cudnn.deterministic = True
+torch.backends.cudnn.benchmark = True
+np.random.seed(seed)
+random.seed(seed)
+
+
+try:  # relative import
+    from model import Model
+    from dataset import BinaryClassifierDataset as Dataset
+    from dataset import get_optimize_class
+except ImportError:
+    from .model import Model
+    from .dataset import BinaryClassifierDataset as Dataset
+    from .dataset import get_optimize_class
+
+# import
+import torch.nn as nn
+from torch import optim
+from torch.optim import lr_scheduler
+from torch.utils.data import DataLoader
+from torch.nn import functional as F
+import os
+import sys
+import warnings
+warnings.filterwarnings("ignore", category=UserWarning)
+
+# load additional config
+import json
+config_file = os.path.join(os.path.dirname(os.path.dirname(__file__)), "config.json")
+with open(config_file, "r") as f:
+    additional_config = json.load(f)
+
+
+
+
+# config
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+config = {
+    "dataset_root": "from_additional_config",
+    "batch_size": 500 if __name__ == "__main__" else 50,
+    "num_workers": 16,
+    "pre_learning_rate": 0.01,
+    "learning_rate": 2e-5,
+    "pre_epochs": 0,
+    "epochs": 50,
+    "weight_decay": 0.1,
+    "save_learning_rate": 1e-6,
+    "total_save_number": 5,
+    "tag": os.path.basename(os.path.dirname(__file__)),
+    "optimize_class": get_optimize_class()[0],
+    "optimize_class_int": get_optimize_class()[1],
+}
+config.update(additional_config)
+print("Training:", config["optimize_class"])
+
+
+
+
+# Data
+dataset = Dataset(
+    root=config["dataset_root"],
+    train=True,
+    optimize_class=config["optimize_class"],
+)
+train_loader = DataLoader(
+    dataset=dataset,
+    batch_size=config["batch_size"],
+    num_workers=config["num_workers"],
+    shuffle=True,
+    drop_last=True,
+    pin_memory=True,
+    persistent_workers=True,
+)
+test_loader = DataLoader(
+    dataset=Dataset(
+        root=config["dataset_root"],
+        train=False,
+        optimize_class=config["optimize_class"],
+    ),
+    batch_size=config["batch_size"],
+    num_workers=config["num_workers"],
+    shuffle=False,
+)
+
+# Model
+model, head = Model()
+model.load_state_dict(torch.load(sys.argv[1], map_location="cpu", weights_only=True))
+model = model.to(device)
+class FocalLoss(nn.Module):
+    def __init__(self, weight=None, gamma=2):
+        super(FocalLoss, self).__init__()
+        self.weight = weight
+        self.gamma = gamma
+    def forward(self, input, target):
+        ce_loss = F.cross_entropy(input, target, reduction='none', weight=self.weight)
+        pt = torch.exp(-ce_loss)
+        focal_loss = (1 - pt) ** self.gamma * ce_loss
+        return focal_loss.mean()
+criterion = FocalLoss()
+
+# Optimizer
+head_optimizer = optim.AdamW(
+    head.parameters(),
+    lr=config["pre_learning_rate"],
+    weight_decay=config["weight_decay"],
+)
+optimizer = optim.AdamW(
+    model.parameters(),
+    lr=config["learning_rate"],
+    weight_decay=config["weight_decay"],
+)
+scheduler = lr_scheduler.CosineAnnealingLR(
+    optimizer,
+    T_max=config["epochs"],
+    eta_min=config["save_learning_rate"],
+)
+
+
+
+
+# Training
+def train(model=model, optimizer=optimizer, scheduler=scheduler):
+    model.train()
+    for batch_idx, (inputs, targets) in enumerate(train_loader):
+        inputs, targets = inputs.to(device), targets.to(device)
+        optimizer.zero_grad()
+        with torch.cuda.amp.autocast(enabled=False, dtype=torch.bfloat16):
+            outputs = model(inputs)
+            loss = criterion(outputs, targets)
+        loss.backward()
+        optimizer.step()
+    if scheduler is not None:
+        scheduler.step()
+
+# test
+@torch.no_grad()
+def test(model=model):
+    model.eval()
+    all_targets = []
+    all_predicts = []
+    test_loss = 0
+    correct = 0
+    total = 0
+    for batch_idx, (inputs, targets) in enumerate(test_loader):
+        inputs, targets = inputs.to(device), targets.to(device)
+        with torch.cuda.amp.autocast(enabled=False, dtype=torch.bfloat16):
+            outputs = model(inputs)
+            loss = criterion(outputs, targets)
+        # to logging losses
+        all_targets.extend(targets.flatten().tolist())
+        test_loss += loss.item()
+        _, predicts = outputs.max(1)
+        all_predicts.extend(predicts.flatten().tolist())
+        total += targets.size(0)
+        correct += predicts.eq(targets).sum().item()
+    loss = test_loss / (batch_idx + 1)
+    acc = correct / total
+    print(f"Loss: {loss:.4f} | Acc: {acc:.4f}\n")
+    model.train()
+    return loss, acc, all_targets, all_predicts
+
+# save train
+def save_train(model=model, optimizer=optimizer):
+    data_loader = DataLoader(
+        dataset=dataset,
+        batch_size=min(len(dataset) // config["total_save_number"], config["batch_size"]),
+        num_workers=config["num_workers"],
+        shuffle=True,
+        drop_last=True,
+    )
+    model.train()
+    for batch_idx, (inputs, targets) in enumerate(data_loader):
+        inputs, targets = inputs.to(device), targets.to(device)
+        optimizer.zero_grad()
+        with torch.cuda.amp.autocast(enabled=False, dtype=torch.bfloat16):
+            outputs = model(inputs)
+            loss = criterion(outputs, targets)
+        loss.backward()
+        optimizer.step()
+        # Save checkpoint
+        # _, acc, _, _ = test(model=model)
+        acc = 1.0
+        if not os.path.isdir('checkpoint'):
+            os.mkdir('checkpoint')
+        save_state = {key: value.cpu().to(torch.float32) for key, value in model.state_dict().items()}
+        torch.save(save_state, f"checkpoint/{str(batch_idx).zfill(4)}_acc{acc:.4f}_class{config['optimize_class_int']}_{config['tag']}.pth")
+        print("save:", f"checkpoint/{str(batch_idx).zfill(4)}_acc{acc:.4f}_class{config['optimize_class_int']}_{config['tag']}.pth")
+        # exit loop
+        if batch_idx+1 == config["total_save_number"]:
+            break
+
+
+
+
+# main
+if __name__ == '__main__':
+    test(model=model)
+    for epoch in range(config["pre_epochs"]):
+        train(model=model, optimizer=head_optimizer, scheduler=None)
+        test(model=model)
+    for epoch in range(config["epochs"]):
+        train(model=model, optimizer=optimizer, scheduler=scheduler)
+        test(model=model)
+    # save_train(model=model, optimizer=optimizer)
+print("time stamp:", time.time())

dataset/condition_classinput_vittiny/model.py

@@ -0,0 +1,25 @@
+import torch
+import torch.nn as nn
+import timm
+
+
+def Model():
+    model = timm.create_model("vit_tiny_patch16_224", pretrained=True)
+    model.head = nn.Sequential(
+        nn.Linear(192, 192, bias=True),
+        nn.SiLU(),
+        nn.Linear(192, 2, bias=False),
+    )
+    for param in model.head.parameters():
+        param = nn.Parameter(torch.ones_like(param) / 192)
+        param.requires_grad = True
+    return model, model.head
+
+
+if __name__ == "__main__":
+    model, _ = Model()
+    print(model)
+    num_param = 0
+    for v in model.parameters():
+        num_param += v.numel()
+    print("num_param:", num_param)

dataset/condition_classinput_vittiny/split.sh

@@ -0,0 +1,28 @@
+mkdir checkpoint_test
+mkdir checkpoint_train
+mkdir generated
+
+mv ./checkpoint/*class0314* ./checkpoint_test
+mv ./checkpoint/*class0482* ./checkpoint_test
+mv ./checkpoint/*class0589* ./checkpoint_test
+mv ./checkpoint/*class0197* ./checkpoint_test
+mv ./checkpoint/*class0462* ./checkpoint_test
+mv ./checkpoint/*class0111* ./checkpoint_test
+mv ./checkpoint/*class0101* ./checkpoint_test
+mv ./checkpoint/*class0278* ./checkpoint_test
+mv ./checkpoint/*class0793* ./checkpoint_test
+mv ./checkpoint/*class0279* ./checkpoint_test
+mv ./checkpoint/*class0653* ./checkpoint_test
+mv ./checkpoint/*class0238* ./checkpoint_test
+mv ./checkpoint/*class1001* ./checkpoint_test
+mv ./checkpoint/*class0141* ./checkpoint_test
+mv ./checkpoint/*class0884* ./checkpoint_test
+mv ./checkpoint/*class0592* ./checkpoint_test
+mv ./checkpoint/*class0502* ./checkpoint_test
+mv ./checkpoint/*class0643* ./checkpoint_test
+mv ./checkpoint/*class0383* ./checkpoint_test
+mv ./checkpoint/*class0128* ./checkpoint_test
+
+mv ./checkpoint/* ./checkpoint_train
+
+rm checkpoint -r

dataset/condition_classinput_vittiny/test.py

@@ -0,0 +1,30 @@
+import os
+import sys
+if __name__ == "__main__":
+    from train import *
+else:  # relative import
+    from .train import *
+
+
+
+
+try:
+    test_item = sys.argv[1]
+except IndexError:
+    assert __name__ == "__main__"
+    test_item = "./checkpoint_test"
+test_items = []
+if os.path.isdir(test_item):
+    for item in os.listdir(test_item):
+        item = os.path.join(test_item, item)
+        test_items.append(item)
+elif os.path.isfile(test_item):
+    test_items.append(test_item)
+
+
+
+
+for item in test_items:
+    state = torch.load(item, map_location="cpu")
+    model.load_state_dict({key: value.to(torch.float32).to(device) for key, value in state.items()})
+    loss, acc, all_targets, all_predicts = test(model=model)

dataset/condition_classinput_vittiny/train.py

@@ -0,0 +1,212 @@
+# set global seed
+import time
+print("time stamp:", time.time())
+import random
+import numpy as np
+import torch
+seed = SEED = 20
+torch.manual_seed(seed)
+torch.cuda.manual_seed(seed)
+torch.cuda.manual_seed_all(seed)
+torch.backends.cudnn.deterministic = True
+torch.backends.cudnn.benchmark = True
+np.random.seed(seed)
+random.seed(seed)
+
+
+try:  # relative import
+    from model import Model
+    from dataset import BinaryClassifierDataset as Dataset
+    from dataset import get_optimize_class
+except ImportError:
+    from .model import Model
+    from .dataset import BinaryClassifierDataset as Dataset
+    from .dataset import get_optimize_class
+
+# import
+import torch.nn as nn
+from torch import optim
+from torch.optim import lr_scheduler
+from torch.utils.data import DataLoader
+from torch.nn import functional as F
+import os
+import sys
+import warnings
+warnings.filterwarnings("ignore", category=UserWarning)
+
+# load additional config
+import json
+config_file = os.path.join(os.path.dirname(os.path.dirname(__file__)), "config.json")
+with open(config_file, "r") as f:
+    additional_config = json.load(f)
+
+
+
+
+# config
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+config = {
+    "dataset_root": "from_additional_config",
+    "batch_size": 500 if __name__ == "__main__" else 50,
+    "num_workers": 16,
+    "pre_learning_rate": 0.01,
+    "learning_rate": 1e-4,
+    "pre_epochs": 2,
+    "epochs": 13,
+    "weight_decay": 0.1,
+    "save_learning_rate": 2e-5,
+    "total_save_number": 5,
+    "tag": os.path.basename(os.path.dirname(__file__)),
+    "optimize_class": get_optimize_class()[0],
+    "optimize_class_int": get_optimize_class()[1],
+}
+config.update(additional_config)
+print("Training:", config["optimize_class"])
+
+
+
+
+# Data
+dataset = Dataset(
+    root=config["dataset_root"],
+    train=True,
+    optimize_class=config["optimize_class"],
+)
+train_loader = DataLoader(
+    dataset=dataset,
+    batch_size=config["batch_size"],
+    num_workers=config["num_workers"],
+    shuffle=True,
+    drop_last=True,
+    pin_memory=True,
+    persistent_workers=True,
+)
+test_loader = DataLoader(
+    dataset=Dataset(
+        root=config["dataset_root"],
+        train=False,
+        optimize_class=config["optimize_class"],
+    ),
+    batch_size=config["batch_size"],
+    num_workers=config["num_workers"],
+    shuffle=False,
+)
+
+# Model
+model, head = Model()
+model = model.to(device)
+class FocalLoss(nn.Module):
+    def __init__(self, weight=None, gamma=2):
+        super(FocalLoss, self).__init__()
+        self.weight = weight
+        self.gamma = gamma
+    def forward(self, input, target):
+        ce_loss = F.cross_entropy(input, target, reduction='none', weight=self.weight)
+        pt = torch.exp(-ce_loss)
+        focal_loss = (1 - pt) ** self.gamma * ce_loss
+        return focal_loss.mean()
+criterion = FocalLoss()
+
+# Optimizer
+head_optimizer = optim.AdamW(
+    head.parameters(),
+    lr=config["pre_learning_rate"],
+    weight_decay=config["weight_decay"],
+)
+optimizer = optim.AdamW(
+    model.parameters(),
+    lr=config["learning_rate"],
+    weight_decay=config["weight_decay"],
+)
+scheduler = lr_scheduler.CosineAnnealingLR(
+    optimizer,
+    T_max=config["epochs"],
+    eta_min=config["save_learning_rate"],
+)
+
+
+
+
+# Training
+def train(model=model, optimizer=optimizer, scheduler=scheduler):
+    model.train()
+    for batch_idx, (inputs, targets) in enumerate(train_loader):
+        inputs, targets = inputs.to(device), targets.to(device)
+        optimizer.zero_grad()
+        with torch.cuda.amp.autocast(enabled=False, dtype=torch.bfloat16):
+            outputs = model(inputs)
+            loss = criterion(outputs, targets)
+        loss.backward()
+        optimizer.step()
+    if scheduler is not None:
+        scheduler.step()
+
+# test
+@torch.no_grad()
+def test(model=model):
+    model.eval()
+    all_targets = []
+    all_predicts = []
+    test_loss = 0
+    correct = 0
+    total = 0
+    for batch_idx, (inputs, targets) in enumerate(test_loader):
+        inputs, targets = inputs.to(device), targets.to(device)
+        with torch.cuda.amp.autocast(enabled=False, dtype=torch.bfloat16):
+            outputs = model(inputs)
+            loss = criterion(outputs, targets)
+        # to logging losses
+        all_targets.extend(targets.flatten().tolist())
+        test_loss += loss.item()
+        _, predicts = outputs.max(1)
+        all_predicts.extend(predicts.flatten().tolist())
+        total += targets.size(0)
+        correct += predicts.eq(targets).sum().item()
+    loss = test_loss / (batch_idx + 1)
+    acc = correct / total
+    print(f"Loss: {loss:.4f} | Acc: {acc:.4f}\n")
+    model.train()
+    return loss, acc, all_targets, all_predicts
+
+# save train
+def save_train(model=model, optimizer=optimizer):
+    data_loader = DataLoader(
+        dataset=dataset,
+        batch_size=min(len(dataset) // config["total_save_number"], config["batch_size"]),
+        num_workers=config["num_workers"],
+        shuffle=True,
+        drop_last=True,
+    )
+    model.train()
+    for batch_idx, (inputs, targets) in enumerate(data_loader):
+        inputs, targets = inputs.to(device), targets.to(device)
+        optimizer.zero_grad()
+        with torch.cuda.amp.autocast(enabled=False, dtype=torch.bfloat16):
+            outputs = model(inputs)
+            loss = criterion(outputs, targets)
+        loss.backward()
+        optimizer.step()
+        # Save checkpoint
+        _, acc, _, _ = test(model=model)
+        if not os.path.isdir('checkpoint'):
+            os.mkdir('checkpoint')
+        save_state = {key: value.cpu().to(torch.float32) for key, value in model.state_dict().items()}
+        torch.save(save_state, f"checkpoint/{str(batch_idx).zfill(4)}_acc{acc:.4f}_class{config['optimize_class_int']}_{config['tag']}.pth")
+        print("save:", f"checkpoint/{str(batch_idx).zfill(4)}_acc{acc:.4f}_class{config['optimize_class_int']}_{config['tag']}.pth")
+        # exit loop
+        if batch_idx+1 == config["total_save_number"]:
+            break
+
+
+
+
+# main
+if __name__ == '__main__':
+    for epoch in range(config["pre_epochs"]):
+        train(model=model, optimizer=head_optimizer, scheduler=None)
+        # test(model=model)
+    for epoch in range(config["epochs"]):
+        train(model=model, optimizer=optimizer, scheduler=scheduler)
+        # test(model=model)
+    save_train(model=model, optimizer=optimizer)
+print("time stamp:", time.time())

dataset/condition_classinput_vittiny/train.sh

@@ -0,0 +1,10 @@
+#!/bin/bash
+
+start=1
+end=1022
+
+for i in $(seq $start $end)
+do
+    python train.py class$i
+    sleep 1
+done

dataset/condition_imageinput_vittiny/README.md

@@ -0,0 +1 @@
+Code for condition_imageinput_vittiny is coming...

dataset/condition_imageinput_vittiny/dataset.py

@@ -0,0 +1,46 @@
+import re
+import sys
+from torch.utils.data import Dataset
+from torchvision.datasets import CIFAR10
+import torchvision.transforms as transforms
+
+
+
+
+class BinaryClassifierDataset(Dataset):
+    def __init__(self, root, train, optimize_class):
+        optimize_class = [optimize_class,] if isinstance(optimize_class, int) else optimize_class
+        self.optimize_class = optimize_class
+        self.dataset = CIFAR10(
+            root=root,
+            train=train,
+            download=True,
+            transform=transforms.Compose([
+                transforms.Resize(224),
+                transforms.RandomHorizontalFlip(),
+                transforms.AutoAugment(policy=transforms.AutoAugmentPolicy("cifar10")),
+                transforms.ToTensor(),
+                transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
+            ])
+        )
+
+    def __getitem__(self, index):
+        img, origin_target = self.dataset[index]
+        target = 1 if origin_target in self.optimize_class else 0
+        return img, target
+
+    def __len__(self):
+        return self.dataset.__len__()
+
+
+
+
+def get_optimize_class():
+    try:  # get string
+        string = sys.argv[1]
+    except IndexError:
+        RuntimeError("sys.argv[1] not found")
+    class_int_string = str(re.search(r'class(\d+)', string).group(1)).zfill(4)
+    one_hot_string = bin(int(class_int_string))[2:].zfill(10)
+    optimize_class = [index for index, i in enumerate(one_hot_string) if i == "1"]
+    return list(optimize_class), class_int_string

dataset/condition_imageinput_vittiny/model.py

@@ -0,0 +1,18 @@
+import torch
+import torch.nn as nn
+import timm
+
+
+def Model():
+    model = timm.create_model("vit_tiny_patch16_224", pretrained=True)
+    model.head = nn.Linear(192, 2)
+    return model, model.head
+
+
+if __name__ == "__main__":
+    model, _ = Model()
+    print(model)
+    num_param = 0
+    for v in model.parameters():
+        num_param += v.numel()
+    print("num_param:", num_param)

dataset/condition_imageinput_vittiny/test.py

@@ -0,0 +1,30 @@
+import os
+import sys
+if __name__ == "__main__":
+    from train import *
+else:  # relative import
+    from .train import *
+
+
+
+
+try:
+    test_item = sys.argv[1]
+except IndexError:
+    assert __name__ == "__main__"
+    test_item = "./checkpoint"
+test_items = []
+if os.path.isdir(test_item):
+    for item in os.listdir(test_item):
+        item = os.path.join(test_item, item)
+        test_items.append(item)
+elif os.path.isfile(test_item):
+    test_items.append(test_item)
+
+
+
+
+for item in test_items:
+    state = torch.load(item, map_location="cpu")
+    model.load_state_dict({key: value.to(torch.float32).to(device) for key, value in state.items()})
+    loss, acc, all_targets, all_predicts = test(model=model)

dataset/condition_imageinput_vittiny/train.py

@@ -0,0 +1,208 @@
+# set global seed
+import random
+import numpy as np
+import torch
+seed = SEED = 20
+torch.manual_seed(seed)
+torch.cuda.manual_seed(seed)
+torch.cuda.manual_seed_all(seed)
+torch.backends.cudnn.deterministic = True
+torch.backends.cudnn.benchmark = True
+np.random.seed(seed)
+random.seed(seed)
+
+try:  # relative import
+    from model import Model
+    from dataset import BinaryClassifierDataset as Dataset
+    from dataset import get_optimize_class
+except ImportError:
+    from .model import Model
+    from .dataset import BinaryClassifierDataset as Dataset
+    from .dataset import get_optimize_class
+
+# import
+import torch.nn as nn
+from torch import optim
+from torch.optim import lr_scheduler
+from torch.utils.data import DataLoader
+from torch.nn import functional as F
+import os
+import sys
+import warnings
+warnings.filterwarnings("ignore", category=UserWarning)
+
+# load additional config
+import json
+config_file = os.path.join(os.path.dirname(os.path.dirname(__file__)), "config.json")
+with open(config_file, "r") as f:
+    additional_config = json.load(f)
+
+
+
+
+# config
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+config = {
+    "dataset_root": "from_additional_config",
+    "batch_size": 250 if __name__ == "__main__" else 50,
+    "num_workers": 20,
+    "pre_learning_rate": 0.01,
+    "learning_rate": 3e-5,
+    "pre_epochs": 2,
+    "epochs": 13,
+    "weight_decay": 0.1,
+    "save_learning_rate": 1e-5,
+    "total_save_number": 10,
+    "tag": os.path.basename(os.path.dirname(__file__)),
+    "optimize_class": get_optimize_class()[0],
+    "optimize_class_int": get_optimize_class()[1],
+}
+config.update(additional_config)
+print("Training:", config["optimize_class"])
+
+
+
+
+# Data
+dataset = Dataset(
+    root=config["dataset_root"],
+    train=True,
+    optimize_class=config["optimize_class"],
+)
+train_loader = DataLoader(
+    dataset=dataset,
+    batch_size=config["batch_size"],
+    num_workers=config["num_workers"],
+    shuffle=True,
+    drop_last=True,
+    pin_memory=True,
+    persistent_workers=True,
+)
+test_loader = DataLoader(
+    dataset=Dataset(
+        root=config["dataset_root"],
+        train=False,
+        optimize_class=config["optimize_class"],
+    ),
+    batch_size=config["batch_size"],
+    num_workers=config["num_workers"],
+    shuffle=False,
+)
+
+# Model
+model, head = Model()
+model = model.to(device)
+class FocalLoss(nn.Module):
+    def __init__(self, weight=None, gamma=2):
+        super(FocalLoss, self).__init__()
+        self.weight = weight
+        self.gamma = gamma
+    def forward(self, input, target):
+        ce_loss = F.cross_entropy(input, target, reduction='none', weight=self.weight)
+        pt = torch.exp(-ce_loss)
+        focal_loss = (1 - pt) ** self.gamma * ce_loss
+        return focal_loss.mean()
+criterion = FocalLoss()
+
+# Optimizer
+head_optimizer = optim.AdamW(
+    head.parameters(),
+    lr=config["pre_learning_rate"],
+    weight_decay=config["weight_decay"],
+)
+optimizer = optim.AdamW(
+    model.parameters(),
+    lr=config["learning_rate"],
+    weight_decay=config["weight_decay"],
+)
+scheduler = lr_scheduler.CosineAnnealingLR(
+    optimizer,
+    T_max=config["epochs"],
+    eta_min=config["save_learning_rate"],
+)
+
+
+
+
+# Training
+def train(model=model, optimizer=optimizer, scheduler=scheduler):
+    model.train()
+    for batch_idx, (inputs, targets) in enumerate(train_loader):
+        inputs, targets = inputs.to(device), targets.to(device)
+        optimizer.zero_grad()
+        with torch.cuda.amp.autocast(enabled=False, dtype=torch.bfloat16):
+            outputs = model(inputs)
+            loss = criterion(outputs, targets)
+        loss.backward()
+        optimizer.step()
+    if scheduler is not None:
+        scheduler.step()
+
+# test
+@torch.no_grad()
+def test(model=model):
+    model.eval()
+    all_targets = []
+    all_predicts = []
+    test_loss = 0
+    correct = 0
+    total = 0
+    for batch_idx, (inputs, targets) in enumerate(test_loader):
+        inputs, targets = inputs.to(device), targets.to(device)
+        with torch.cuda.amp.autocast(enabled=False, dtype=torch.bfloat16):
+            outputs = model(inputs)
+            loss = criterion(outputs, targets)
+        # to logging losses
+        all_targets.extend(targets.flatten().tolist())
+        test_loss += loss.item()
+        _, predicts = outputs.max(1)
+        all_predicts.extend(predicts.flatten().tolist())
+        total += targets.size(0)
+        correct += predicts.eq(targets).sum().item()
+    loss = test_loss / (batch_idx + 1)
+    acc = correct / total
+    print(f"Loss: {loss:.4f} | Acc: {acc:.4f}\n")
+    model.train()
+    return loss, acc, all_targets, all_predicts
+
+# save train
+def save_train(model=model, optimizer=optimizer):
+    data_loader = DataLoader(
+        dataset=dataset,
+        batch_size=min(len(dataset) // config["total_save_number"], config["batch_size"]),
+        num_workers=config["num_workers"],
+        shuffle=True,
+        drop_last=True,
+    )
+    model.train()
+    for batch_idx, (inputs, targets) in enumerate(data_loader):
+        inputs, targets = inputs.to(device), targets.to(device)
+        optimizer.zero_grad()
+        with torch.cuda.amp.autocast(enabled=False, dtype=torch.bfloat16):
+            outputs = model(inputs)
+            loss = criterion(outputs, targets)
+        loss.backward()
+        optimizer.step()
+        # Save checkpoint
+        _, acc, _, _ = test(model=model)
+        if not os.path.isdir('checkpoint'):
+            os.mkdir('checkpoint')
+        save_state = {key: value.cpu().to(torch.float32) for key, value in model.state_dict().items()}
+        torch.save(save_state, f"checkpoint/{str(batch_idx).zfill(4)}_acc{acc:.4f}_class{config['optimize_class_int']}_{config['tag']}.pth")
+        print("save:", f"checkpoint/{str(batch_idx).zfill(4)}_acc{acc:.4f}_class{config['optimize_class_int']}_{config['tag']}.pth")
+        # exit loop
+        if batch_idx+1 == config["total_save_number"]:
+            break
+
+
+
+
+# main
+if __name__ == '__main__':
+    for epoch in range(config["pre_epochs"]):
+        train(model=model, optimizer=head_optimizer, scheduler=None)
+        test(model=model)
+    for epoch in range(config["epochs"]):
+        train(model=model, optimizer=optimizer, scheduler=scheduler)
+        test(model=model)
+    save_train(model=model, optimizer=optimizer)

dataset/condition_imageinput_vittiny/train.sh

@@ -0,0 +1,11 @@
+#!/bin/bash
+
+start=0
+end=9
+
+for i in $(seq $start $end)
+do
+    power=$((2**i))
+    CUDA_VISIBLE_DEVICES=5 python train.py class$power
+    sleep 1
+done

dataset/condition_permutation_vittiny/model.py

@@ -0,0 +1,18 @@
+import torch
+import torch.nn as nn
+import timm
+
+
+def Model():
+    model = timm.create_model("vit_tiny_patch16_224", pretrained=False)
+    model.head = nn.Linear(192, 10)
+    return model, model.head
+
+
+if __name__ == "__main__":
+    model, _ = Model()
+    print(model)
+    num_param = 0
+    for v in model.parameters():
+        num_param += v.numel()
+    print("num_param:", num_param)

dataset/condition_permutation_vittiny/test.py

@@ -0,0 +1,31 @@
+import os
+import sys
+if __name__ == "__main__":
+    from train import *
+else:  # relative import
+    from .train import *
+
+
+
+
+try:
+    test_item = sys.argv[1]
+except IndexError:
+    assert __name__ == "__main__"
+    test_item = "./checkpoint"
+test_items = []
+if os.path.isdir(test_item):
+    for item in os.listdir(test_item):
+        item = os.path.join(test_item, item)
+        test_items.append(item)
+elif os.path.isfile(test_item):
+    test_items.append(test_item)
+
+
+
+
+for item in test_items:
+    print(f"testing: {item}")
+    state = torch.load(item, map_location="cpu")
+    model.load_state_dict({key: value.to(torch.float32).to(device) for key, value in state.items()})
+    loss, acc, all_targets, all_predicts = test(model=model)

dataset/condition_permutation_vittiny/train.py

@@ -0,0 +1,210 @@
+# set global seed
+import time
+print("time stamp:", time.time())
+import random
+import numpy as np
+import torch
+import re
+import sys
+if __name__ == "__main__":
+    def get_permutation_state():
+        try:  # get string
+            string = sys.argv[1]
+        except IndexError:
+            RuntimeError("sys.argv[1] not found")
+        class_int_string = str(re.search(r'class(\d+)', string).group(1)).zfill(4)
+        return int(class_int_string)
+    seed = SEED = get_permutation_state()
+else:  # when testing
+    seed = SEED = 0
+torch.manual_seed(seed)
+torch.cuda.manual_seed(seed)
+torch.cuda.manual_seed_all(seed)
+torch.backends.cudnn.deterministic = True
+torch.backends.cudnn.benchmark = True
+np.random.seed(seed)
+random.seed(seed)
+print("Seed:", SEED)
+
+try:  # relative import
+    from model import Model
+except ImportError:
+    from .model import Model
+
+# import
+import torch.nn as nn
+from torch import optim
+from torch.optim import lr_scheduler
+from torch.utils.data import DataLoader
+from torchvision.datasets import CIFAR10 as Dataset
+from torchvision import transforms
+from torch.nn import functional as F
+import warnings
+warnings.filterwarnings("ignore", category=UserWarning)
+
+# load additional config
+import os
+import json
+config_file = os.path.join(os.path.dirname(os.path.dirname(__file__)), "config.json")
+with open(config_file, "r") as f:
+    additional_config = json.load(f)
+
+
+
+
+# config
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+config = {
+    "dataset_root": "from_additional_config",
+    "batch_size": 250 if __name__ == "__main__" else 50,
+    "num_workers": 16,
+    "learning_rate": 5e-3,
+    "epochs": 200,
+    "weight_decay": 0.1,
+    "save_learning_rate": 2e-5,
+    "total_save_number": 5,
+    "tag": os.path.basename(os.path.dirname(__file__)),
+}
+config.update(additional_config)
+
+
+
+
+# Data
+dataset = Dataset(
+    root=config["dataset_root"],
+    train=True,
+    download=True,
+    transform=transforms.Compose([
+        transforms.Resize(224),
+        transforms.RandomCrop(224, padding=32),
+        transforms.RandomHorizontalFlip(),
+        transforms.AutoAugment(policy=transforms.AutoAugmentPolicy("cifar10")),
+        transforms.ToTensor(),
+        transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2471, 0.2435, 0.2616)),
+    ])
+)
+train_loader = DataLoader(
+    dataset=dataset,
+    batch_size=config["batch_size"],
+    num_workers=config["num_workers"],
+    shuffle=True,
+    drop_last=True,
+    pin_memory=True,
+    persistent_workers=True,
+)
+test_loader = DataLoader(
+    dataset=Dataset(
+        root=config["dataset_root"],
+        train=False,
+        download=True,
+        transform=transforms.Compose([
+            transforms.Resize(224),
+            transforms.ToTensor(),
+            transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2471, 0.2435, 0.2616)),
+        ])),
+    batch_size=config["batch_size"],
+    num_workers=config["num_workers"],
+    shuffle=False,
+)
+
+# Model
+model, head = Model()
+model = model.to(device)
+criterion = nn.CrossEntropyLoss()
+
+# Optimizer
+optimizer = optim.AdamW(
+    model.parameters(),
+    lr=config["learning_rate"],
+    weight_decay=config["weight_decay"],
+)
+scheduler = lr_scheduler.CosineAnnealingLR(
+    optimizer,
+    T_max=config["epochs"],
+    eta_min=config["save_learning_rate"],
+)
+
+
+
+
+# Training
+def train(model=model, optimizer=optimizer, scheduler=scheduler):
+    model.train()
+    for batch_idx, (inputs, targets) in enumerate(train_loader):
+        inputs, targets = inputs.to(device), targets.to(device)
+        optimizer.zero_grad()
+        with torch.cuda.amp.autocast(enabled=True, dtype=torch.bfloat16):
+            outputs = model(inputs)
+            loss = criterion(outputs, targets)
+        loss.backward()
+        optimizer.step()
+    if scheduler is not None:
+        scheduler.step()
+
+# test
+@torch.no_grad()
+def test(model=model):
+    model.eval()
+    all_targets = []
+    all_predicts = []
+    test_loss = 0
+    correct = 0
+    total = 0
+    for batch_idx, (inputs, targets) in enumerate(test_loader):
+        inputs, targets = inputs.to(device), targets.to(device)
+        with torch.cuda.amp.autocast(enabled=True, dtype=torch.bfloat16):
+            outputs = model(inputs)
+            loss = criterion(outputs, targets)
+        # to logging losses
+        all_targets.extend(targets.flatten().tolist())
+        test_loss += loss.item()
+        _, predicts = outputs.max(1)
+        all_predicts.extend(predicts.flatten().tolist())
+        total += targets.size(0)
+        correct += predicts.eq(targets).sum().item()
+    loss = test_loss / (batch_idx + 1)
+    acc = correct / total
+    print(f"Loss: {loss:.4f} | Acc: {acc:.4f}\n")
+    model.train()
+    return loss, acc, all_targets, all_predicts
+
+# save train
+def save_train(model=model, optimizer=optimizer):
+    data_loader = DataLoader(
+        dataset=dataset,
+        batch_size=min(len(dataset) // config["total_save_number"], config["batch_size"]),
+        num_workers=config["num_workers"],
+        shuffle=True,
+        drop_last=True,
+    )
+    model.train()
+    for batch_idx, (inputs, targets) in enumerate(data_loader):
+        inputs, targets = inputs.to(device), targets.to(device)
+        optimizer.zero_grad()
+        with torch.cuda.amp.autocast(enabled=False, dtype=torch.bfloat16):
+            outputs = model(inputs)
+            loss = criterion(outputs, targets)
+        loss.backward()
+        optimizer.step()
+        # Save checkpoint
+        _, acc, _, _ = test(model=model)
+        if not os.path.isdir('checkpoint'):
+            os.mkdir('checkpoint')
+        save_state = {key: value.cpu().to(torch.float32) for key, value in model.state_dict().items()}
+        torch.save(save_state, f"checkpoint/{str(batch_idx).zfill(4)}_acc{acc:.4f}_class{SEED:04d}_{config['tag']}.pth")
+        print("save:", f"checkpoint/{str(batch_idx).zfill(4)}_acc{acc:.4f}_class{SEED:04d}_{config['tag']}.pth")
+        # exit loop
+        if batch_idx+1 == config["total_save_number"]:
+            break
+
+
+
+
+# main
+if __name__ == '__main__':
+    for epoch in range(config["epochs"]):
+        train(model=model, optimizer=optimizer, scheduler=scheduler)
+        test(model=model)
+    save_train(model=model, optimizer=optimizer)
+print("time stamp:", time.time())

dataset/condition_permutation_vittiny/train.sh

@@ -0,0 +1,10 @@
+#!/bin/bash
+
+start=0
+end=19
+
+for i in $(seq $start $end)
+do
+    python train.py class$i
+    sleep 1
+done

dataset/config.json

@@ -0,0 +1 @@
+{"dataset_root": "path_to_your_dataset", "imagenet_root": {"train": null, "test": null}, "dora_root": "/home/wangkai/arpgen/DoRA/commonsense_reasoning", "dora_env_name": "dora_llama"}

dataset/dataset.py

@@ -0,0 +1,327 @@
+import torch
+import einops
+from torch.utils.data import Dataset
+from torchvision.datasets import CIFAR10
+from torchvision import transforms
+import os
+import math
+import random
+import json
+from abc import ABC
+import pickle
+
+
+
+
+def pad_to_length(x, common_factor, **config):
+    if x.numel() % common_factor == 0:
+        return x.flatten()
+    # print(f"padding {x.shape} according to {common_factor}")
+    full_length = (x.numel() // common_factor + 1) * common_factor
+    padding_length = full_length - len(x.flatten())
+    padding = torch.full([padding_length, ], dtype=x.dtype, device=x.device, fill_value=config["fill_value"])
+    x = torch.cat((x.flatten(), padding), dim=0)
+    return x
+
+def layer_to_token(x, common_factor, **config):
+    if config["granularity"] == 2:  # split by output
+        if x.numel() <= common_factor:
+            return pad_to_length(x.flatten(), common_factor, **config)[None]
+        dim2 = x[0].numel()
+        dim1 = x.shape[0]
+        if dim2 <= common_factor:
+            i = int(dim1 / (common_factor / dim2))
+            while True:
+                if dim1 % i == 0 and dim2 * (dim1 // i) <= common_factor:
+                    output = x.view(-1, dim2 * (dim1 // i))
+                    output = [pad_to_length(item, common_factor, **config) for item in output]
+                    return torch.stack(output, dim=0)
+                i += 1
+        else:  # dim2 > common_factor
+            output = [layer_to_token(item, common_factor, **config) for item in x]
+            return torch.cat(output, dim=0)
+    elif config["granularity"] == 1:  # split by layer
+        return pad_to_length(x.flatten(), common_factor, **config).view(-1, common_factor)
+    elif config["granularity"] == 0:  # flatten directly
+        return x.flatten()
+    else:  # NotImplementedError
+        raise NotImplementedError("granularity: 0: flatten directly, 1: split by layer, 2: split by output dim")
+
+
+def token_to_layer(tokens, shape, **config):
+    common_factor = tokens.shape[-1]
+    if config["granularity"] == 2:  # split by output
+        num_element = math.prod(shape)
+        if num_element <= common_factor:
+            param = tokens[0][:num_element].view(shape)
+            tokens = tokens[1:]
+            return param, tokens
+        dim2 = num_element // shape[0]
+        dim1 = shape[0]
+        if dim2 <= common_factor:
+            i = int(dim1 / (common_factor / dim2))
+            while True:
+                if dim1 % i == 0 and dim2 * (dim1 // i) <= common_factor:
+                    item_per_token = dim2 * (dim1 // i)
+                    length = num_element // item_per_token
+                    output = [item[:item_per_token] for item in tokens[:length]]
+                    param = torch.cat(output, dim=0).view(shape)
+                    tokens = tokens[length:]
+                    return param, tokens
+                i += 1
+        else:  # dim2 > common_factor
+            output = []
+            for i in range(shape[0]):
+                param, tokens = token_to_layer(tokens, shape[1:], **config)
+                output.append(param.flatten())
+            param = torch.cat(output, dim=0).view(shape)
+            return param, tokens
+    elif config["granularity"] == 1:  # split by layer
+        num_element = math.prod(shape)
+        token_num = num_element // common_factor if num_element % common_factor == 0 \
+                else num_element // common_factor + 1
+        param = tokens.flatten()[:num_element].view(shape)
+        tokens = tokens[token_num:]
+        return param, tokens
+    elif config["granularity"] == 0:  # flatten directly
+        num_element = math.prod(shape)
+        param = tokens.flatten()[:num_element].view(shape)
+        tokens = pad_to_length(tokens.flatten()[num_element:],
+                common_factor, fill_value=torch.nan).view(-1, common_factor)
+        return param, tokens
+    else:  # NotImplementedError
+        raise NotImplementedError("granularity: 0: flatten directly, 1: split by layer, 2: split by output dim")
+
+
+def positional_embedding_2d(dim1, dim2, d_model):
+    assert d_model % 4 == 0, f"Cannot use sin/cos positional encoding with odd dimension {d_model}"
+    pe = torch.zeros(d_model, dim1, dim2)
+    d_model = int(d_model / 2)  # Each dimension use half of d_model
+    div_term = torch.exp(torch.arange(0., d_model, 2, dtype=torch.float32) * -(math.log(10000.0) / d_model))
+    pos_w = torch.arange(0., dim2).unsqueeze(1)
+    pos_h = torch.arange(0., dim1).unsqueeze(1)
+    pe[0:d_model:2, :, :] = torch.sin(pos_w * div_term).transpose(0, 1).unsqueeze(1).repeat(1, dim1, 1)
+    pe[1:d_model:2, :, :] = torch.cos(pos_w * div_term).transpose(0, 1).unsqueeze(1).repeat(1, dim1, 1)
+    pe[d_model::2, :, :] = torch.sin(pos_h * div_term).transpose(0, 1).unsqueeze(2).repeat(1, 1, dim2)
+    pe[d_model+1::2, :, :] = torch.cos(pos_h * div_term).transpose(0, 1).unsqueeze(2).repeat(1, 1, dim2)
+    return pe.permute(1, 2, 0)
+
+
+def positional_embedding_1d(dim1, d_model):
+    pe = torch.zeros(dim1, d_model)
+    position = torch.arange(0, dim1, dtype=torch.float).unsqueeze(1)
+    div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
+    pe[:, 0::2] = torch.sin(position * div_term)
+    pe[:, 1::2] = torch.cos(position * div_term)
+    return pe
+
+
+
+
+class BaseDataset(Dataset, ABC):
+    data_path = None
+    generated_path = None
+    test_command = None
+    config = {
+        "fill_value": torch.nan,
+        "granularity": 1,  # 0: flatten directly, 1: split by layer, 2: split by output
+        "pe_granularity": 2,  # 0: no embedding, 1: 1d embedding, 2: 2d embedding
+    }
+
+    def __init__(self, checkpoint_path=None, dim_per_token=8192, **kwargs):
+        if not os.path.exists(self.data_path):
+            os.makedirs(self.data_path, exist_ok=False)
+        if self.generated_path is not None and not os.path.exists(os.path.dirname(self.generated_path)):
+            os.makedirs(os.path.dirname(self.generated_path))
+        self.config.update(kwargs)
+        checkpoint_path = self.data_path if checkpoint_path is None else checkpoint_path
+        assert os.path.exists(checkpoint_path)
+        self.dim_per_token = dim_per_token
+        self.structure = None  # set in get_structure()
+        self.sequence_length = None  # set in get_structure()
+        # load checkpoint_list
+        checkpoint_list = os.listdir(checkpoint_path)
+        self.checkpoint_list = list([os.path.join(checkpoint_path, item) for item in checkpoint_list])
+        self.length = self.real_length = len(self.checkpoint_list)
+        self.set_infinite_dataset()
+        # get structure
+        structure_cache_file = os.path.join(os.path.dirname(self.data_path), "structure.cache")
+        try:  # try to load cache file
+            assert os.path.exists(structure_cache_file)
+            with open(structure_cache_file, "rb") as f:
+                print(f"Loading cache from {structure_cache_file}")
+                cache_file = pickle.load(f)
+            if len(self.checkpoint_list) != 0:
+                assert set(cache_file["checkpoint_list"]) == set(self.checkpoint_list)
+                self.structure = cache_file["structure"]
+            else:  # empty checkpoint_list, only generate
+                print("Cannot find any trained checkpoint, loading cache file for generating!")
+                self.structure = cache_file["structure"]
+                fake_diction = {key: torch.zeros(item[0]) for key, item in self.structure.items()}
+                torch.save(fake_diction, os.path.join(checkpoint_path, "fake_checkpoint.pth"))
+                self.checkpoint_list.append(os.path.join(checkpoint_path, "fake_checkpoint.pth"))
+                self.length = self.real_length = len(self.checkpoint_list)
+                self.set_infinite_dataset()
+                os.system(f"rm {os.path.join(checkpoint_path, 'fake_checkpoint.pth')}")
+        except AssertionError:  # recompute cache file
+            print("==> Organizing structure..")
+            self.structure = self.get_structure()
+            with open(structure_cache_file, "wb") as f:
+                pickle.dump({"structure": self.structure, "checkpoint_list": self.checkpoint_list}, f)
+        # get sequence_length
+        self.sequence_length = self.get_sequence_length()
+
+    def get_sequence_length(self):
+        fake_diction = {key: torch.zeros(item[0]) for key, item in self.structure.items()}
+        # get sequence_length
+        param = self.preprocess(fake_diction)
+        self.sequence_length = param.size(0)
+        return self.sequence_length
+
+    def get_structure(self):
+        # get structure
+        checkpoint_list = self.checkpoint_list
+        structures = [{} for _ in range(len(checkpoint_list))]
+        for i, checkpoint in enumerate(checkpoint_list):
+            diction = torch.load(checkpoint, map_location="cpu")
+            for key, value in diction.items():
+                if ("num_batches_tracked" in key) or (value.numel() == 1) or not torch.is_floating_point(value):
+                    structures[i][key] = (value.shape, value, None)
+                elif "running_var" in key:
+                    pre_mean = value.mean() * 0.95
+                    value = torch.log(value / pre_mean + 0.05)
+                    structures[i][key] = (value.shape, pre_mean, value.mean(), value.std())
+                else:  # conv & linear
+                    structures[i][key] = (value.shape, value.mean(), value.std())
+        final_structure = {}
+        structure_diction = torch.load(checkpoint_list[0], map_location="cpu")
+        for key, param in structure_diction.items():
+            if ("num_batches_tracked" in key) or (param.numel() == 1) or not torch.is_floating_point(param):
+                final_structure[key] = (param.shape, param, None)
+            elif "running_var" in key:
+                value = [param.shape, 0., 0., 0.]
+                for structure in structures:
+                    for i in [1, 2, 3]:
+                        value[i] += structure[key][i]
+                for i in [1, 2, 3]:
+                    value[i] /= len(structures)
+                final_structure[key] = tuple(value)
+            else:  # conv & linear
+                value = [param.shape, 0., 0.]
+                for structure in structures:
+                    for i in [1, 2]:
+                        value[i] += structure[key][i]
+                for i in [1, 2]:
+                    value[i] /= len(structures)
+                final_structure[key] = tuple(value)
+        self.structure = final_structure
+        return self.structure
+
+    def set_infinite_dataset(self, max_num=None):
+        if max_num is None:
+            max_num = self.length * 1000000
+        self.length = max_num
+        return self
+
+    @property
+    def max_permutation_state(self):
+        return self.real_length
+
+    def get_position_embedding(self, positional_embedding_dim=None):
+        if positional_embedding_dim is None:
+            positional_embedding_dim = self.dim_per_token // 2
+        assert self.structure is not None, "run get_structure before get_position_embedding"
+        if self.config["pe_granularity"] == 2:
+            print("Use 2d positional embedding")
+            positional_embedding_index = []
+            for key, item in self.structure.items():
+                if ("num_batches_tracked" in key) or (item[-1] is None):
+                    continue
+                else:  # conv & linear
+                    shape, *_ = item
+                fake_param = torch.ones(size=shape)
+                fake_param = layer_to_token(fake_param, self.dim_per_token, **self.config)
+                positional_embedding_index.append(list(range(fake_param.size(0))))
+            dim1 = len(positional_embedding_index)
+            dim2 = max([len(token_per_layer) for token_per_layer in positional_embedding_index])
+            full_pe = positional_embedding_2d(dim1, dim2, positional_embedding_dim)
+            positional_embedding = []
+            for layer_index, token_indexes in enumerate(positional_embedding_index):
+                for token_index in token_indexes:
+                    this_pe = full_pe[layer_index, token_index]
+                    positional_embedding.append(this_pe)
+            positional_embedding = torch.stack(positional_embedding)
+            return positional_embedding
+        elif self.config["pe_granularity"] == 1:
+            print("Use 1d positional embedding")
+            return positional_embedding_1d(self.sequence_length, positional_embedding_dim)
+        elif self.config["pe_granularity"] == 0:
+            print("Not use positional embedding")
+            return torch.zeros_like(self.__getitem__(0))
+        else:  # NotImplementedError
+            raise NotImplementedError("pe_granularity: 0: no embedding, 1: 1d embedding, 2: 2d embedding")
+
+    def __len__(self):
+        return self.length
+
+    def __getitem__(self, index):
+        index = index % self.real_length
+        diction = torch.load(self.checkpoint_list[index], map_location="cpu")
+        param = self.preprocess(diction)
+        return param, index
+
+    def save_params(self, params, save_path):
+        diction = self.postprocess(params.cpu().to(torch.float32))
+        torch.save(diction, save_path)
+
+    def preprocess(self, diction: dict, **kwargs) -> torch.Tensor:
+        param_list = []
+        for key, value in diction.items():
+            if ("num_batches_tracked" in key) or (value.numel() == 1) or not torch.is_floating_point(value):
+                continue
+            elif "running_var" in key:
+                shape, pre_mean, mean, std = self.structure[key]
+                value = torch.log(value / pre_mean + 0.05)
+            else:  # normal
+                shape, mean, std = self.structure[key]
+            value = (value - mean) / std
+            value = layer_to_token(value, self.dim_per_token, **self.config)
+            param_list.append(value)
+        param = torch.cat(param_list, dim=0)
+        if self.config["granularity"] == 0:  # padding directly process tail
+            param = pad_to_length(param, self.dim_per_token, **self.config).view(-1, self.dim_per_token)
+        # print("Sequence length:", param.size(0))
+        return param.to(torch.float32)
+
+    def postprocess(self, params: torch.Tensor, **kwargs) -> dict:
+        diction = {}
+        params = params if len(params.shape) == 2 else params.squeeze(0)
+        for key, item in self.structure.items():
+            if ("num_batches_tracked" in key) or (item[-1] is None):
+                shape, mean, std = item
+                diction[key] = mean
+                continue
+            elif "running_var" in key:
+                shape, pre_mean, mean, std = item
+            else:  # conv & linear
+                shape, mean, std = item
+            this_param, params = token_to_layer(params, shape, **self.config)
+            this_param = this_param * std + mean
+            if "running_var" in key:
+                this_param = torch.clip(torch.exp(this_param) - 0.05, min=0.001) * pre_mean
+            diction[key] = this_param
+        return diction
+
+
+class ConditionalDataset(BaseDataset, ABC):
+    def _extract_condition(self, index: int):
+        name = self.checkpoint_list[index]
+        condition_list = os.path.basename(name).split("_")
+        return condition_list
+
+    def __getitem__(self, index):
+        index = index % self.real_length
+        diction = torch.load(self.checkpoint_list[index], map_location="cpu")
+        condition = self._extract_condition(index)
+        param = self.preprocess(diction)
+        return param, condition

dataset/downtask_detection/README.md

@@ -0,0 +1 @@
+Code for segmentation is coming...

dataset/downtask_detection/test.sh

@@ -0,0 +1,11 @@
+#!/usr/bin/env bash
+
+source /path/to/miniconda3/bin/activate /path/to/miniconda3/envs/environment
+
+CLUSTER=True \
+DETECTRON2_DATASETS="/path/to/" \
+PYTHONPATH="$(dirname $0)/Detection":$PYTHONPATH \
+python $(dirname $0)/Detection/tools/lazyconfig_train_net.py --config-file $(dirname $0)/Detection/projects/ViTDet/configs/COCO/our_vit_b_100ep.py --finetune "VIT_BASE_IN21K" \
+--num-gpus 1 \
+--fulltune \
+--eval-only "train.init_checkpoint='$1'"

dataset/downtask_dora_r16/adapter_config.json

@@ -0,0 +1,23 @@
+{
+  "Wdecompose_target_modules": null,
+  "base_model_name_or_path": "yahma/llama-7b-hf",
+  "bias": "none",
+  "dora_simple": true,
+  "enable_lora": null,
+  "fan_in_fan_out": false,
+  "inference_mode": true,
+  "lora_alpha": 32,
+  "lora_dropout": 0.05,
+  "merge_weights": false,
+  "modules_to_save": null,
+  "peft_type": "DORA",
+  "r": 16,
+  "target_modules": [
+    "q_proj",
+    "k_proj",
+    "v_proj",
+    "up_proj",
+    "down_proj"
+  ],
+  "task_type": "CAUSAL_LM"
+}