AlekseySh
commented
4 months ago hey, @olegkorshunov and thank you for using OML!
Yes, it does work with ArcFace on GPU. You simply forgot to move your criterion to the cuda:0, which is needed because ArcFace has its classification head.
Here is a fully working example (the one from Readme, but with ArcFace):
from torch.optim import Adam
from torch.utils.data import DataLoader
from oml import datasets as d
from oml.inference import inference
from oml.losses import ArcFaceLoss
from oml.metrics import calc_retrieval_metrics_rr
from oml.models import ViTExtractor
from oml.registry import get_transforms_for_pretrained
from oml.retrieval import RetrievalResults, AdaptiveThresholding
from oml.samplers import BalanceSampler
from oml.utils import get_mock_images_dataset
device = "cuda:0"
model = ViTExtractor.from_pretrained("vits16_dino").to(device).train()
transform, _ = get_transforms_for_pretrained("vits16_dino")
df_train, df_val = get_mock_images_dataset(global_paths=True)
train = d.ImageLabeledDataset(df_train, transform=transform)
val = d.ImageQueryGalleryLabeledDataset(df_val, transform=transform)
optimizer = Adam(model.parameters(), lr=1e-4)
criterion = ArcFaceLoss(in_features=384, num_classes=4).to(device) # <<<<<<<<<<<<<<<<<< HERE
sampler = BalanceSampler(train.get_labels(), n_labels=2, n_instances=2)
def training():
for batch in DataLoader(train, batch_sampler=sampler):
embeddings = model(batch["input_tensors"].to(device))
loss = criterion(embeddings, batch["labels"].to(device))
loss.backward()
optimizer.step()
optimizer.zero_grad()
print(criterion.last_logs)
def validation():
embeddings = inference(model, val, batch_size=4, num_workers=0)
rr = RetrievalResults.from_embeddings(embeddings, val, n_items=3)
rr = AdaptiveThresholding(n_std=2).process(rr)
rr.visualize(query_ids=[2, 1], dataset=val, show=True)
print(calc_retrieval_metrics_rr(rr, map_top_k=(3,), cmc_top_k=(1,)))
training()
validation()
olegkorshunov
error log