import pandas as pd
import numpy as np
from tqdm import tqdm
from sklearn.model_selection import train_test_split
import torchmetrics
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from torchvision import models, transforms
from torchvision.transforms import InterpolationMode
seed = 42
root_path = "/home/stefan/ioai-prep/kits/neoai/terminal_animals"
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
torch.manual_seed(seed)
<torch._C.Generator at 0x7fc4a273d1d0>
Data
train_df = pd.read_csv(f"{root_path}/train.csv")
test_df = pd.read_csv(f"{root_path}/test.csv")
train_df.head()
img label
0 \\/\\\\?—\\ccccc/ccccco\P/c;.|o\ooPPPPPPP/PPPP... 32
1 \P/P|\\\ooo||o—||\\@/■\■■■\\■■——\\■■\||■\@\\/|... 32
2 \\\\\\/—\—/c/|\|c\cccc/\c\/|\/c\c/|cc—/—/\//\/... 34
3 —/—////??@@■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■... 24
4 \\/\/\\///\\/\/\|c\|o/o||/|\\..../c\ ... 19
tr, ev = train_test_split(train_df, test_size=0.15, random_state=seed, stratify=train_df["label"])
tr.shape, ev.shape
((3128, 2), (552, 2))
all_text = "".join(train_df["img"].tolist())
vocab = sorted(list(set(all_text)))
char_to_idx = {char: i + 1 for i, char in enumerate(vocab)}
char_to_idx[""] = 0
vocab_size = len(char_to_idx)
vocab_size
16
def text_to_idx(img_str):
lines = img_str.strip().split("\n")
grid = np.zeros((128, 128), dtype=np.int64)
for i, line in enumerate(lines[:128]):
line_indices = [char_to_idx.get(c, 0) for c in line[:128]]
grid[i, : len(line_indices)] = line_indices
return grid
class ASCIIDataset(Dataset):
def __init__(self, df, transform=None, is_test=False):
self.df = df
self.transform = transform
self.is_test = is_test
def __len__(self):
return len(self.df)
def __getitem__(self, idx):
grid = text_to_idx(self.df.iloc[idx]["img"])
img_tensor = torch.from_numpy(grid).unsqueeze(0).to(torch.float32)
if self.transform:
img_tensor = self.transform(img_tensor)
img_tensor = img_tensor.squeeze(0).to(torch.long)
if self.is_test:
return img_tensor
return img_tensor, torch.tensor(self.df.iloc[idx]["label"], dtype=torch.long)
num_classes = train_df["label"].nunique()
train_transforms = transforms.Compose([
transforms.RandomAffine(degrees=0, translate=(0.1, 0.1), interpolation=InterpolationMode.NEAREST),
transforms.RandomHorizontalFlip(),
])
train_loader = DataLoader(
ASCIIDataset(tr, transform=train_transforms), batch_size=64, shuffle=True
)
val_loader = DataLoader(ASCIIDataset(ev), batch_size=64, shuffle=False)
Model
class ResNetASCII(nn.Module):
def __init__(self, vocab_size, num_classes, emb_dim=64):
super().__init__()
self.embedding = nn.Embedding(vocab_size, emb_dim)
self.stem = nn.Sequential(
nn.Conv2d(emb_dim, 64, 3, 1, 1, bias=False),
nn.BatchNorm2d(64),
nn.ReLU(inplace=True),
nn.MaxPool2d(3, 2, 1),
)
self.resnet = models.resnet34(weights=None)
self.resnet.conv1 = self.resnet.bn1 = self.resnet.relu = self.resnet.maxpool = nn.Identity()
self.resnet.fc = nn.Sequential(
nn.Dropout(0.5),
nn.Linear(self.resnet.fc.in_features, num_classes)
)
def forward(self, x):
x = self.embedding(x).permute(0, 3, 1, 2)
return self.resnet(self.stem(x))
Pretraining
class ConvMAE(nn.Module):
def __init__(self, vocab_size, emb_dim=128, mask_ratio=0.75):
super().__init__()
self.mask_ratio = mask_ratio
self.emb_dim = emb_dim
self.embedding = nn.Embedding(vocab_size, emb_dim, padding_idx=0)
self.mask_token = nn.Parameter(torch.randn(emb_dim) * 0.02)
self.encoder = nn.Sequential(
nn.Conv2d(emb_dim, 128, 3, stride=2, padding=1),
nn.BatchNorm2d(128),
nn.GELU(),
nn.Conv2d(128, 256, 3, stride=2, padding=1),
nn.BatchNorm2d(256),
nn.GELU(),
nn.Conv2d(256, 256, 3, padding=1),
nn.BatchNorm2d(256),
nn.GELU(),
)
self.decoder = nn.Sequential(
nn.ConvTranspose2d(256, 128, 4, stride=2, padding=1),
nn.BatchNorm2d(128),
nn.GELU(),
nn.ConvTranspose2d(128, emb_dim, 4, stride=2, padding=1),
nn.BatchNorm2d(emb_dim),
nn.GELU(),
nn.Conv2d(emb_dim, vocab_size, 1),
)
def forward(self, x):
B, H, W = x.shape
mask = torch.rand(B, H, W, device=x.device) < self.mask_ratio
emb = self.embedding(x)
mask_token = self.mask_token.view(1, 1, 1, self.emb_dim).expand(
B, H, W, self.emb_dim
)
emb = torch.where(mask.unsqueeze(-1), mask_token, emb)
emb = emb.permute(0, 3, 1, 2)
encoded = self.encoder(emb)
logits = self.decoder(encoded)
return logits, mask
pretrain_loader = DataLoader(ASCIIDataset(train_df), batch_size=64, shuffle=True)
mae = ConvMAE(vocab_size, emb_dim=64, mask_ratio=0.75).to(device)
mae_criterion = nn.CrossEntropyLoss(reduction="none")
mae_epochs = 50
mae_optim = optim.AdamW(mae.parameters(), lr=1e-3, weight_decay=0.05)
mae_scheduler = optim.lr_scheduler.OneCycleLR(
mae_optim,
max_lr=2e-3,
epochs=mae_epochs,
steps_per_epoch=len(pretrain_loader),
pct_start=0.05,
)
for epoch in range(1, mae_epochs + 1):
mae.train()
total_loss = 0
for imgs, _ in tqdm(pretrain_loader, desc=f"MAE Epoch {epoch}"):
imgs = imgs.to(device)
mae_optim.zero_grad()
logits, mask = mae(imgs)
loss = mae_criterion(logits, imgs)
weights = torch.where(mask & (imgs != 0),
torch.tensor(4.0, device=device),
torch.tensor(1.0, device=device))
weights = weights * (imgs != 0).float()
loss = (loss * weights).sum() / weights.sum().clamp(min=1)
loss.backward()
nn.utils.clip_grad_norm_(mae.parameters(), 1.0)
mae_optim.step()
mae_scheduler.step()
total_loss += loss.item()
print(f"MAE Epoch {epoch:3d} | Loss: {total_loss/len(pretrain_loader):.4f} | LR: {mae_scheduler.get_last_lr()[0]:.2e}")
torch.save(mae.embedding.state_dict(), f"{root_path}/mae_embedding.pth")
model = ResNetASCII(vocab_size, num_classes).to(device)
model.embedding.load_state_dict(torch.load(f"{root_path}/mae_embedding.pth"))
model.embedding.requires_grad_(False)
Embedding(16, 64)
Training
best_f1 = 0
epochs = 100
patience = 20
no_improve = 0
f1_metric = torchmetrics.F1Score(task="multiclass", num_classes=num_classes, average="micro").to(device)
optimizer = optim.AdamW(model.parameters(), lr=2e-3)
criterion = nn.CrossEntropyLoss(label_smoothing=0.1)
scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)
for epoch in range(1, epochs+1):
model.train()
train_loss = 0
for imgs, labels in tqdm(train_loader):
imgs, labels = imgs.to(device), labels.to(device)
optimizer.zero_grad()
loss = criterion(model(imgs), labels)
loss.backward()
optimizer.step()
train_loss += loss.item()
model.eval()
with torch.no_grad():
for imgs, labels in val_loader:
pred = model(imgs.to(device)).argmax(1)
f1_metric.update(pred, labels.to(device))
val_f1 = f1_metric.compute()
f1_metric.reset()
scheduler.step()
if val_f1 > best_f1:
best_f1 = val_f1
no_improve = 0
torch.save(model.state_dict(), f"{root_path}/best_v2.pth")
print(f"New Best F1: {val_f1:.4f}")
else:
no_improve += 1
print(f"Epoch {epoch} | Loss: {train_loss/len(train_loader):.4f} | F1: {val_f1:.4f}")
if no_improve >= patience:
print(f"Early stopping at epoch {epoch}")
break
Submission
model.load_state_dict(torch.load(f"{root_path}/best_v2.pth"))
test_loader = DataLoader(ASCIIDataset(test_df, is_test=True), batch_size=48)
model.eval()
final_preds = []
with torch.no_grad():
for imgs in tqdm(test_loader):
out = model(imgs.to(device))
final_preds.extend(out.argmax(1).cpu().numpy())
100%|██████████| 77/77 [00:07<00:00, 10.92it/s]
test_df["label"] = final_preds
test_df.to_csv(f"{root_path}/submission.csv", columns=["id", "label"], index=None)