dnn_face_recognition to Pytorch(,ONNX,CoreML,...)

[pfeatherstone]

template <template <int,template<typename>class,int,typename> class block, int N, template<typename>class BN, typename SUBNET>
using residual = add_prev1<block<N,BN,1,tag1<SUBNET>>>;

template <template <int,template<typename>class,int,typename> class block, int N, template<typename>class BN, typename SUBNET>
using residual_down = add_prev2<avg_pool<2,2,2,2,skip1<tag2<block<N,BN,2,tag1<SUBNET>>>>>>;

template <int N, template <typename> class BN, int stride, typename SUBNET> 
using block  = BN<con<N,3,3,1,1,relu<BN<con<N,3,3,stride,stride,SUBNET>>>>>;

template <int N, typename SUBNET> using ares      = relu<residual<block,N,affine,SUBNET>>;
template <int N, typename SUBNET> using ares_down = relu<residual_down<block,N,affine,SUBNET>>;

template <typename SUBNET> using alevel0 = ares_down<256,SUBNET>;
template <typename SUBNET> using alevel1 = ares<256,ares<256,ares_down<256,SUBNET>>>;
template <typename SUBNET> using alevel2 = ares<128,ares<128,ares_down<128,SUBNET>>>;
template <typename SUBNET> using alevel3 = ares<64,ares<64,ares<64,ares_down<64,SUBNET>>>>;
template <typename SUBNET> using alevel4 = ares<32,ares<32,ares<32,SUBNET>>>;

using anet_type = loss_metric<fc_no_bias<128,avg_pool_everything<
                            alevel0<
                            alevel1<
                            alevel2<
                            alevel3<
                            alevel4<
                            max_pool<3,3,2,2,relu<affine<con<32,7,7,2,2,
                            input_rgb_image_sized<150>
                            >>>>>>>>>>>>;

anet_type net;
deserialize("dlib_face_recognition_resnet_model_v1.dat") >> net;

struct visitor
{
    std::ofstream out;

    visitor(const char* weights_path)
    : out(weights_path, std::ios::binary)
    {
    }

    template <typename T> 
    void operator()(size_t layer_index, T& layer) {}

    template <long nf, long nr, long nc, int sy, int sx, int py, int px, typename SUBNET>
    void operator()(size_t layer_index, add_layer<con_<nf, nr, nc, sy, sx, py, px>, SUBNET>& l)
    {
        tensor& params = l.layer_details().get_layer_params();
        DLIB_CASSERT(!l.layer_details().bias_is_disabled(), "bias shouldn't be disabled in this model");

        const long nf_in    = (params.size() - nf) / (nf*nr*nc);
        auto filters        = alias_tensor(nf, nf_in, nr, nc);
        auto biases         = alias_tensor(1, nf);
        auto f              = filters(params, 0);
        auto b              = biases(params, filters.size());
        out.write((const char*)f.host(), f.size()*sizeof(float));
        out.write((const char*)b.host(), b.size()*sizeof(float));
    }

    template <typename SUBNET> void operator()(size_t layer_index, add_layer<affine_, SUBNET>& l)
    {
        if (!l.layer_details().is_disabled())
        {
            auto f = l.layer_details().get_gamma();
            auto b = l.layer_details().get_beta();
            out.write((const char*)f.host(), f.size()*sizeof(float));
            out.write((const char*)b.host(), b.size()*sizeof(float));
        }
    }

    // fully connected layers
    template <unsigned long num_outputs, typename SUBNET>
    void operator()(size_t layer_index, add_layer<fc_<num_outputs, FC_NO_BIAS>, SUBNET>& l)
    {
        auto weight = l.layer_details().get_weights();
        matrix<float> tmp = trans(mat(weight));
        out.write((const char*)&tmp(0,0), tmp.size()*sizeof(float));
    }
};

visit_layers_backwards(net, visitor("dnn_face_recognition.bin"));

Then:

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from   einops import rearrange

class InputRGBImage(nn.Module):
    def __init__(self):
        super().__init__()
        self.avg = nn.Parameter(rearrange(torch.tensor([122.782, 117.001, 104.298]), 'c -> () c () ()'), requires_grad=False)
    def forward(self, x):
        x = rearrange(x, 'b h w c -> b c h w').float()
        x = (x - self.avg) / 256.0
        return x

class Residual(nn.Module):
    def __init__(self, module, stride=1):
        super().__init__()
        self.module = module
        self.stride = stride
    def forward(self, x):
        r        = x if self.stride == 1 else F.avg_pool2d(x, kernel_size=2, stride=2, padding=0)
        y        = self.module(x)
        C, H, W  = max(r.shape[1], y.shape[1]), max(r.shape[2], y.shape[2]), max(r.shape[3], y.shape[3])
        r       = F.pad(r, (0,W-r.shape[3],0,H-r.shape[2],0,C-r.shape[1]))
        y       = F.pad(y, (0,W-y.shape[3],0,H-y.shape[2],0,C-y.shape[1]))
        return r + y

class GlobalAveragePooling(nn.Module):
    def forward(self, x):
        return x.mean((2,3))
    
def block(c1, c2, stride=1):
    pad = 1 if stride==1 else 0
    return nn.Sequential(Residual(nn.Sequential(nn.Conv2d(c1,c2,kernel_size=3,stride=stride,padding=pad),
                                                nn.BatchNorm2d(c2),
                                                nn.ReLU(True),
                                                nn.Conv2d(c2,c2,3,1,1),
                                                nn.BatchNorm2d(c2)), stride),
                         nn.ReLU(True))

def DlibFaceRecognition():
    return nn.Sequential(InputRGBImage(),
                         nn.Conv2d(3,32,kernel_size=7,stride=2,padding=0),
                         nn.BatchNorm2d(32),
                         nn.ReLU(True),
                         nn.MaxPool2d(kernel_size=3, stride=2, padding=0),
                         block(32,32),      block(32,32), block(32,32),
                         block(32, 64, 2),  block(64,64), block(64,64), block(64,64),
                         block(64, 128, 2), block(128,128), block(128,128),
                         block(128, 256, 2), block(256,256), block(256,256),
                         block(256, 256, 2),
                         GlobalAveragePooling(),
                         nn.Linear(256, 128, bias=False))

n = DlibFaceRecognition().eval()

with open('dnn_face_recognition.bin', mode='rb') as file:
    for m in n.modules():
        if isinstance(m, nn.Conv2d) or isinstance(m, nn.BatchNorm2d) or isinstance(m, nn.Linear):
            if m.weight is not None:
                w = np.fromfile(file, dtype=np.float32, count=m.weight.numel()).reshape(m.weight.shape)
                m.weight.data.copy_(torch.from_numpy(w))
            if m.bias is not None:
                b = np.fromfile(file, dtype=np.float32, count=m.bias.numel()).reshape(m.bias.shape)
                m.bias.data.copy_(torch.from_numpy(b))

# Try it
x = torch.randint(0, 256, (32, 150,150,3))
y = n(x)

[pfeatherstone]

Let's make a "Tiny" version. Who can reduce the most number of lines without sacrificing readability.

[pfeatherstone]

con_ could do with .get_weights() and .get_biases() member functions...

[pfeatherstone]

We could also do with a visit_computational_layers_backwards() function. I checked the docs and couldn't find one.

[pfeatherstone]

Also, add_prev is oddly permissive when it comes to tensor shapes. they don't need to match.

[pfeatherstone]

@davisking this could be an interesting example to add to the examples folder. With some code added to export to ONNX and CoreML.

[davisking]

Yeah that would make for a cool example. Totally send me a PR :)

[pfeatherstone]

OK will do at some point.

[pfeatherstone]

It would be cool if we could save dlib weights to .npy file. Looking at https://github.com/llohse/libnpy, it looks pretty easy.

[pfeatherstone]

Then we could do something like:

anet_type net;
deserialize("dlib_face_recognition_resnet_model_v1.dat") >> net;
save_weights_tonpy(net, "dlib_face_recognition_resnet_model_v1.npy")

Or maybe something like https://github.com/huggingface/safetensors

[davisking]

Yeah. It's only going to work sometimes though. Since all these things are a little bit different to each other.

[pfeatherstone]

Yeah it wouldn't be full-proof just slightly better than saving to raw binary file