手撕代码insightFace中的arcface_torch

背景介绍

　　搞人脸识别的同学基本都听过insightFace 的大名，在开源工程里面可以帮助大伙快速的建立自己的baseline , 代码玩儿的溜的同学说不一定一两天就玩儿通了．原始的insightface是mxnet实现的，但是现在工业界和学术的有非常多的人使用pytorch作为自己的开发平台，这就带来了一定的局部不适．终于最近insightFace基于pytorch实现了相关的算法，并公布了性能还不错的开源预训练模型，pytorch建议到1.6.0版本以上．感谢相关的开源作者做出的贡献．

　　那么这个工程能给你带来什么呢？

　　１．可以迅速建立一个baseline.如果是小公司没有人力去研究新的算法，设计新的loss，或者模型，根据吴恩达大神的80%数据20％算法原则，就无脑持续收集数据，不断训练模型就可以了．简单吧，愉快吧．

　　２．优秀的代码风格，还有简洁的算法实现，如果要拿他作为一个基准然后持续的研究新想法，新loss设计，新模型修改，都是一个不错的验证平台，各种对标的验证集，训练集，应有尽有，开箱即用，方便！

整体介绍：

　　说了这么多，废话就没了，来介绍一下整体工程，介于本篇文章着眼点在于代码分析，理论部分就尽量少一些，而且网上讲理论的文章太多了，同学你就自行度娘吧．工程主体结构如下：

.
├── README.md
├── backbones
│   ├── __init__.py
│   └── iresnet.py      # backbone定义
├── config.py　　　　　　 # 配置文件，使用什么数据集训练，验证，batchsize，学习率等
├── dataset.py          # dataloader 的实现，并行加速异步缓存，
├── docs
│   ├── eval.md　　　　　 # 
│   ├── install.md      # 
│   └── modelzoo.md     # 
├── eval
│   ├── __init__.py
│   └── verification.py
├── eval_ijbc.py
├── kill_all_distributed_python.sh
├── losses.py
├── partial_fc.py
├── requirement.txt
├── run.sh
├── train.py
├── trainNode0.sh
├── trainNode1.sh
├── trainNode2.sh
├── trainSingleNode.sh
└── utils
    ├── __init__.py
    ├── plot.py
    ├── utils_amp.py    # 混合精度
    ├── utils_callbacks.py
    ├── utils_logging.py
    └── utils_os.py

我加了一些中文注解，方便童鞋门的理解，并快速开始玩儿自己的项目传送门如下：

https://github.com/leoluopy/pytorch_arcface_cosface_partialFC

有相关问题搜索知识星球号：1453755 【CV老司机】加入星球提问。扫码也可加入：

也可以搜索关注微信公众号: CV老司机

相关代码和详细资源或者相关问题，可联系牛先生小猪wx号: jishudashou

后面是一些详细的讲解．

数据加载dataloader：

dataloader其实是一个非常基础的训练组件，在数据量少的时候，怎么写都可以，也没啥区别。不过一旦数据量大了之后，比如达到千万级别后，硬盘IO的读写效率，数据预处理并行度都成为性能瓶颈。

下面的代码即是为了解决上面的问题，一种异步加载，个人认为是比较漂亮的写法。同时，这样封装后，dataloader的遍历方式并没有发生改变，这是十分可贵的一点。

主要的思路是： 1. Thread 内部自动加载数据至queue 2. 重载了 _iter 和 next_ 使得上层的接口保持不变。

主流程和关键点还有一个数据集的加速，平时我们的数据也就几万，多的时候几十万，或者几百万，人脸数据集目前最好的都达到千万级别如glint360k,就有36W的ID 和 1600W样子的图片，这么多的数据，对于磁盘访问，CPU和GPU的均衡是一个考验。同时为了做到pytorch上层接口的保持不变，insight face 做了一个挺漂亮的封装，重载了 _iter, _next_ ， 并在dataloader的内部开起新线程异步不断加载数据到Queue , 外部使用相同接口迭代数据时，就直接从Queue中取，而不需要取得时候才着手数据前处理。迭代完成后，数据触发一个 StopIteration的信号结束迭代。

其他地方还有什么疑问，或者有什么探讨的地方，进知识星球提问吧。

class BackgroundGenerator(threading.Thread):
    def __init__(self, generator, local_rank, max_prefetch=6):
        super(BackgroundGenerator, self).__init__()
        self.queue = Queue.Queue(max_prefetch)
        self.generator = generator
        self.local_rank = local_rank
        self.daemon = True
        self.start()

    def run(self):
        torch.cuda.set_device(self.local_rank)
        for item in self.generator:
            self.queue.put(item)
        self.queue.put(None)

    def next(self):
        next_item = self.queue.get()
        if next_item is None:
            raise StopIteration
        return next_item

    def __next__(self):
        return self.next()

    def __iter__(self):
        return self


class DataLoaderX(DataLoader):
    def __init__(self, local_rank, **kwargs):
        super(DataLoaderX, self).__init__(**kwargs)
        self.stream = torch.cuda.Stream(local_rank)
        self.local_rank = local_rank

    def __iter__(self):
        self.iter = super(DataLoaderX, self).__iter__()
        self.iter = BackgroundGenerator(self.iter, self.local_rank)
        self.preload()
        return self

    def preload(self):
        self.batch = next(self.iter, None)
        if self.batch is None:
            return None
        with torch.cuda.stream(self.stream):
            for k in range(len(self.batch)):
                self.batch[k] = self.batch[k].to(device