|
| 1 | +import argparse |
| 2 | +import os |
| 3 | +import os.path as osp |
| 4 | + |
| 5 | +import torch |
| 6 | +import torch.nn.functional as F |
| 7 | +from PIL import Image |
| 8 | +from torch import Tensor |
| 9 | +from tqdm import tqdm |
| 10 | +from transformers import AutoImageProcessor, AutoModel |
| 11 | + |
| 12 | +from torch_frame import stype |
| 13 | +from torch_frame.config import ImageEmbedder, ImageEmbedderConfig |
| 14 | +from torch_frame.data import DataLoader |
| 15 | +from torch_frame.datasets import DiamondImages |
| 16 | +from torch_frame.nn import ( |
| 17 | + EmbeddingEncoder, |
| 18 | + FTTransformer, |
| 19 | + LinearEmbeddingEncoder, |
| 20 | + LinearEncoder, |
| 21 | +) |
| 22 | + |
| 23 | +parser = argparse.ArgumentParser() |
| 24 | +parser.add_argument("--channels", type=int, default=256) |
| 25 | +parser.add_argument("--num_layers", type=int, default=4) |
| 26 | +parser.add_argument("--batch_size", type=int, default=512) |
| 27 | +parser.add_argument("--lr", type=float, default=0.0001) |
| 28 | +parser.add_argument("--epochs", type=int, default=30) |
| 29 | +parser.add_argument("--seed", type=int, default=0) |
| 30 | +parser.add_argument( |
| 31 | + "--model", |
| 32 | + type=str, |
| 33 | + default="google/vit-base-patch16-224-in21k", |
| 34 | + choices=[ |
| 35 | + "microsoft/resnet-18", |
| 36 | + "google/vit-base-patch16-224-in21k", |
| 37 | + "microsoft/swin-base-patch4-window7-224-in22k", |
| 38 | + ], |
| 39 | +) |
| 40 | + |
| 41 | +args = parser.parse_args() |
| 42 | + |
| 43 | +# Image Embedded |
| 44 | +# ================ ResNet =================== |
| 45 | +# Best Val Acc: 0.2864, Best Test Acc: 0.2789 |
| 46 | +# ================== ViT ==================== |
| 47 | +# Best Val Acc: 0.4173, Best Test Acc: 0.4110 |
| 48 | +# ================= Swin ==================== |
| 49 | +# Best Val Acc: 0.4345, Best Test Acc: 0.4274 |
| 50 | + |
| 51 | + |
| 52 | +class ImageToEmbedding(ImageEmbedder): |
| 53 | + def __init__(self, model_name: str, device: torch.device): |
| 54 | + super().__init__() |
| 55 | + self.model_name = model_name |
| 56 | + self.preprocess = AutoImageProcessor.from_pretrained(model_name) |
| 57 | + self.model = AutoModel.from_pretrained(model_name).to(device) |
| 58 | + self.model.eval() |
| 59 | + self.device = device |
| 60 | + |
| 61 | + def forward_embed(self, images: list[Image]) -> Tensor: |
| 62 | + inputs = self.preprocess(images, return_tensors="pt") |
| 63 | + inputs["pixel_values"] = inputs["pixel_values"].to(self.device) |
| 64 | + with torch.no_grad(): |
| 65 | + res = self.model(**inputs).pooler_output.cpu().detach() |
| 66 | + if "resnet" in self.model_name: |
| 67 | + res = res.squeeze(dim=(2, 3)) |
| 68 | + return res |
| 69 | + |
| 70 | + |
| 71 | +torch.manual_seed(args.seed) |
| 72 | +device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| 73 | + |
| 74 | +# Prepare datasets |
| 75 | +path = osp.join(osp.dirname(osp.realpath(__file__)), "..", "data", |
| 76 | + "diamond_images") |
| 77 | +os.makedirs(path, exist_ok=True) |
| 78 | + |
| 79 | +col_to_image_embedder_cfg = ImageEmbedderConfig( |
| 80 | + image_embedder=ImageToEmbedding(args.model, device), batch_size=10) |
| 81 | +dataset = DiamondImages(path, |
| 82 | + col_to_image_embedder_cfg=col_to_image_embedder_cfg) |
| 83 | + |
| 84 | +model_name = args.model.replace('/', '') |
| 85 | +filename = f"{model_name}_data.pt" |
| 86 | +dataset.materialize(path=osp.join(path, filename)) |
| 87 | +dataset = dataset.shuffle() |
| 88 | +train_dataset, val_dataset, test_dataset = dataset[:0.8], dataset[ |
| 89 | + 0.8:0.9], dataset[0.9:] |
| 90 | + |
| 91 | +train_tensor_frame = train_dataset.tensor_frame |
| 92 | +val_tensor_frame = val_dataset.tensor_frame |
| 93 | +test_tensor_frame = test_dataset.tensor_frame |
| 94 | +train_loader = DataLoader(train_tensor_frame, batch_size=args.batch_size, |
| 95 | + shuffle=True) |
| 96 | +val_loader = DataLoader(val_tensor_frame, batch_size=args.batch_size) |
| 97 | +test_loader = DataLoader(test_tensor_frame, batch_size=args.batch_size) |
| 98 | + |
| 99 | +stype_encoder_dict = { |
| 100 | + stype.categorical: EmbeddingEncoder(), |
| 101 | + stype.numerical: LinearEncoder(), |
| 102 | + stype.image_embedded.parent: LinearEmbeddingEncoder(), |
| 103 | +} |
| 104 | + |
| 105 | +model = FTTransformer( |
| 106 | + channels=args.channels, |
| 107 | + out_channels=dataset.num_classes, |
| 108 | + num_layers=args.num_layers, |
| 109 | + col_stats=dataset.col_stats, |
| 110 | + col_names_dict=train_tensor_frame.col_names_dict, |
| 111 | + stype_encoder_dict=stype_encoder_dict, |
| 112 | +).to(device) |
| 113 | +optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr) |
| 114 | + |
| 115 | + |
| 116 | +def train(epoch: int) -> float: |
| 117 | + model.train() |
| 118 | + loss_accum = total_count = 0 |
| 119 | + |
| 120 | + for tf in tqdm(train_loader, desc=f"Epoch: {epoch}"): |
| 121 | + tf = tf.to(device) |
| 122 | + pred = model(tf) |
| 123 | + loss = F.cross_entropy(pred, tf.y) |
| 124 | + optimizer.zero_grad() |
| 125 | + loss.backward() |
| 126 | + loss_accum += float(loss) * len(tf.y) |
| 127 | + total_count += len(tf.y) |
| 128 | + optimizer.step() |
| 129 | + return loss_accum / total_count |
| 130 | + |
| 131 | + |
| 132 | +@torch.no_grad() |
| 133 | +def test(loader: DataLoader) -> float: |
| 134 | + model.eval() |
| 135 | + accum = total_count = 0 |
| 136 | + |
| 137 | + for tf in loader: |
| 138 | + tf = tf.to(device) |
| 139 | + pred = model(tf) |
| 140 | + pred_class = pred.argmax(dim=-1) |
| 141 | + accum += float((tf.y == pred_class).sum()) |
| 142 | + total_count += len(tf.y) |
| 143 | + |
| 144 | + accuracy = accum / total_count |
| 145 | + return accuracy |
| 146 | + |
| 147 | + |
| 148 | +metric = "Acc" |
| 149 | +best_val_metric = 0 |
| 150 | +best_test_metric = 0 |
| 151 | + |
| 152 | +for epoch in range(1, args.epochs + 1): |
| 153 | + train_loss = train(epoch) |
| 154 | + train_metric = test(train_loader) |
| 155 | + val_metric = test(val_loader) |
| 156 | + test_metric = test(test_loader) |
| 157 | + if val_metric > best_val_metric: |
| 158 | + best_val_metric = val_metric |
| 159 | + best_test_metric = test_metric |
| 160 | + |
| 161 | + print(f"Train Loss: {train_loss:.4f}, Train {metric}: {train_metric:.4f}, " |
| 162 | + f"Val {metric}: {val_metric:.4f}, Test {metric}: {test_metric:.4f}") |
| 163 | + |
| 164 | +print(f"Best Val {metric}: {best_val_metric:.4f}, " |
| 165 | + f"Best Test {metric}: {best_test_metric:.4f}") |
0 commit comments