Extracting features from ArcFace

[saedr]

Hi, I have a question regarding the feature extraction, as I cannot reproduce the results with my own preprocessed files. Given IJB-B-512, your checkpoint for CASIA, and pytorch implementation of ArcFace. I came up with the following code:

import numpy as np
from tqdm import tqdm

import torch
import torch.nn as nn
import torchvision.transforms as transforms
from torchvision import datasets

from model import Backbone

transform = transforms.Compose([
    transforms.RandomHorizontalFlip(),
    transforms.Resize(size=(112, 112)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
    ])

data_path = "../data/IJB-B-512/"
batch_size = 16
num_workers = 16

data = datasets.ImageFolder(data_path, transform=transform)
loader = torch.utils.data.DataLoader(data, 
                                     batch_size=batch_size, 
                                     num_workers=num_workers,
                                     shuffle=True,
                                     pin_memory=True)

model = Backbone(50, 0.6, 'ir_se')
ckpt = torch.load("../pretrained/model_ir_se50.pth")
model.load_state_dict(ckpt)
model.cuda()
model.eval()

features = []
def hook(module, input, output):
    N, C, H, W = output.shape
    output = output.reshape(N, C, -1)
    features.append(output.mean(dim=2).cpu().detach().numpy())

handle = model._modules['body'][23].res_layer[5].fc2.register_forward_hook(hook)
for i_batch, inputs in tqdm(enumerate(loader), total=len(loader)):
    _ = model(inputs[0].cuda())

features = np.concatenate(features)
handle.remove()

Could you please let me know if my approach makes sense or how is it different from yours or could you kindly share your pre-processing module?

[ghost]

Hi Saed (@saedr) , Please how did you make the code run ? Like the code is written for Pytorch 0.4 and Python 2.7 ,and it is difficult to make it run in the new GPU-s. Any suggestion would be helpful.

[houguanqun]

Hi, @saedr , Could you share the dataset "IJB-B"?