Obtain features from ArcFace

vision
1

I’m trying to extract features using pytorch implementation of ArcFace

Here is my 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.ToTensor(),
    transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
    ])

data_path = "../data/CASIA/"
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()

where Backbone model is imported from here, the pre-trained model is downloaded from here, and I’m working with CASIA dataset.

However I get the following error:

Traceback (most recent call last):
  File "preprocess.py", line 86, in <module>
    feat, labels = extract_feature(args.data_path + args.data_str, args.batch_size, args.workers)
  File "preprocess.py", line 67, in extract_feature
    _ = model(inputs[0].cuda())
  File "/home/saed/.local/lib/python3.7/site-packages/torch/nn/modules/module.py", line 532, in __call__
    result = self.forward(*input, **kwargs)
  File "/home/saed/Dropbox/Projects/X_Domain_Clustering/src/arcface.py", line 138, in forward
    x = self.output_layer(x)
  File "/home/saed/.local/lib/python3.7/site-packages/torch/nn/modules/module.py", line 532, in __call__
    result = self.forward(*input, **kwargs)
  File "/home/saed/.local/lib/python3.7/site-packages/torch/nn/modules/container.py", line 100, in forward
    input = module(input)
  File "/home/saed/.local/lib/python3.7/site-packages/torch/nn/modules/module.py", line 532, in __call__
    result = self.forward(*input, **kwargs)
  File "/home/saed/.local/lib/python3.7/site-packages/torch/nn/modules/linear.py", line 87, in forward
    return F.linear(input, self.weight, self.bias)
  File "/home/saed/.local/lib/python3.7/site-packages/torch/nn/functional.py", line 1370, in linear
    ret = torch.addmm(bias, input, weight.t())
RuntimeError: size mismatch, m1: [16 x 131072], m2: [25088 x 512] at /pytorch/aten/src/THC/generic/THCTensorMathBlas.cu:290

Now I’m wondering if my approach is correct, and if so, how should I resolve this issue?

1 Like
2

This error is raised in the output_layer, as the incoming number of features is not matching the defined in_features in this layer.

How did you define the in_features?
If the model was working before, did you change the spatial size of your input?
Also, is your batch size 16? If not, a view or reshape operation might have been called in a wrong way inside your model.

1 Like
3

Thanks for the reply.
I just managed to fix the error by adding the following line in transform:

transforms.Resize(size=(112, 112))

Now my code works. I just wonder if the whole approach makes sense in extracting features using a pre-trained model. The use of register_forward_hook and creating handle, for instance.

1 Like
4

Hello, my name is Azi. from Uzbekistan.My brain is stucked since last week. Please can you explain me what is ArcFace really is? is it loss function?

My real question is : I am writing a code it get ArcFace weights and why I have to use other model’s weights?

Please help me… (uacoding01@gmail.com)