RabbitRUI/ruispace

1import logging2import os3 4import torch5import torch.distributed as dist6from torch.nn import Module7from torch.nn.functional import normalize, linear8from torch.nn.parameter import Parameter9 10 11class PartialFC(Module):12    """13    Author: {Xiang An, Yang Xiao, XuHan Zhu} in DeepGlint,14    Partial FC: Training 10 Million Identities on a Single Machine15    See the original paper:16    https://arxiv.org/abs/2010.0522217    """18 19    @torch.no_grad()20    def __init__(self, rank, local_rank, world_size, batch_size, resume,21                 margin_softmax, num_classes, sample_rate=1.0, embedding_size=512, prefix="./"):22        """23        rank: int24            Unique process(GPU) ID from 0 to world_size - 1.25        local_rank: int26            Unique process(GPU) ID within the server from 0 to 7.27        world_size: int28            Number of GPU.29        batch_size: int30            Batch size on current rank(GPU).31        resume: bool32            Select whether to restore the weight of softmax.33        margin_softmax: callable34            A function of margin softmax, eg: cosface, arcface.35        num_classes: int36            The number of class center storage in current rank(CPU/GPU), usually is total_classes // world_size,37            required.38        sample_rate: float39            The partial fc sampling rate, when the number of classes increases to more than 2 millions, Sampling40            can greatly speed up training, and reduce a lot of GPU memory, default is 1.0.41        embedding_size: int42            The feature dimension, default is 512.43        prefix: str44            Path for save checkpoint, default is './'.45        """46        super(PartialFC, self).__init__()47        #48        self.num_classes: int = num_classes49        self.rank: int = rank50        self.local_rank: int = local_rank51        self.device: torch.device = torch.device("cuda:{}".format(self.local_rank))52        self.world_size: int = world_size53        self.batch_size: int = batch_size54        self.margin_softmax: callable = margin_softmax55        self.sample_rate: float = sample_rate56        self.embedding_size: int = embedding_size57        self.prefix: str = prefix58        self.num_local: int = num_classes // world_size + int(rank < num_classes % world_size)59        self.class_start: int = num_classes // world_size * rank + min(rank, num_classes % world_size)60        self.num_sample: int = int(self.sample_rate * self.num_local)61 62        self.weight_name = os.path.join(self.prefix, "rank_{}_softmax_weight.pt".format(self.rank))63        self.weight_mom_name = os.path.join(self.prefix, "rank_{}_softmax_weight_mom.pt".format(self.rank))64 65        if resume:66            try:67                self.weight: torch.Tensor = torch.load(self.weight_name)68                self.weight_mom: torch.Tensor = torch.load(self.weight_mom_name)69                if self.weight.shape[0] != self.num_local or self.weight_mom.shape[0] != self.num_local:70                    raise IndexError71                logging.info("softmax weight resume successfully!")72                logging.info("softmax weight mom resume successfully!")73            except (FileNotFoundError, KeyError, IndexError):74                self.weight = torch.normal(0, 0.01, (self.num_local, self.embedding_size), device=self.device)75                self.weight_mom: torch.Tensor = torch.zeros_like(self.weight)76                logging.info("softmax weight init!")77                logging.info("softmax weight mom init!")78        else:79            self.weight = torch.normal(0, 0.01, (self.num_local, self.embedding_size), device=self.device)80            self.weight_mom: torch.Tensor = torch.zeros_like(self.weight)81            logging.info("softmax weight init successfully!")82            logging.info("softmax weight mom init successfully!")83        self.stream: torch.cuda.Stream = torch.cuda.Stream(local_rank)84 85        self.index = None86        if int(self.sample_rate) == 1:87            self.update = lambda: 088            self.sub_weight = Parameter(self.weight)89            self.sub_weight_mom = self.weight_mom90        else:91            self.sub_weight = Parameter(torch.empty((0, 0)).cuda(local_rank))92 93    def save_params(self):94        """ Save softmax weight for each rank on prefix95        """96        torch.save(self.weight.data, self.weight_name)97        torch.save(self.weight_mom, self.weight_mom_name)98 99    @torch.no_grad()100    def sample(self, total_label):101        """102        Sample all positive class centers in each rank, and random select neg class centers to filling a fixed103        `num_sample`.104 105        total_label: tensor106            Label after all gather, which cross all GPUs.107        """108        index_positive = (self.class_start <= total_label) & (total_label < self.class_start + self.num_local)109        total_label[~index_positive] = -1110        total_label[index_positive] -= self.class_start111        if int(self.sample_rate) != 1:112            positive = torch.unique(total_label[index_positive], sorted=True)113            if self.num_sample - positive.size(0) >= 0:114                perm = torch.rand(size=[self.num_local], device=self.device)115                perm[positive] = 2.0116                index = torch.topk(perm, k=self.num_sample)[1]117                index = index.sort()[0]118            else:119                index = positive120            self.index = index121            total_label[index_positive] = torch.searchsorted(index, total_label[index_positive])122            self.sub_weight = Parameter(self.weight[index])123            self.sub_weight_mom = self.weight_mom[index]124 125    def forward(self, total_features, norm_weight):126        """ Partial fc forward, `logits = X * sample(W)`127        """128        torch.cuda.current_stream().wait_stream(self.stream)129        logits = linear(total_features, norm_weight)130        return logits131 132    @torch.no_grad()133    def update(self):134        """ Set updated weight and weight_mom to memory bank.135        """136        self.weight_mom[self.index] = self.sub_weight_mom137        self.weight[self.index] = self.sub_weight138 139    def prepare(self, label, optimizer):140        """141        get sampled class centers for cal softmax.142 143        label: tensor144            Label tensor on each rank.145        optimizer: opt146            Optimizer for partial fc, which need to get weight mom.147        """148        with torch.cuda.stream(self.stream):149            total_label = torch.zeros(150                size=[self.batch_size * self.world_size], device=self.device, dtype=torch.long)151            dist.all_gather(list(total_label.chunk(self.world_size, dim=0)), label)152            self.sample(total_label)153            optimizer.state.pop(optimizer.param_groups[-1]['params'][0], None)154            optimizer.param_groups[-1]['params'][0] = self.sub_weight155            optimizer.state[self.sub_weight]['momentum_buffer'] = self.sub_weight_mom156            norm_weight = normalize(self.sub_weight)157            return total_label, norm_weight158 159    def forward_backward(self, label, features, optimizer):160        """161        Partial fc forward and backward with model parallel162 163        label: tensor164            Label tensor on each rank(GPU)165        features: tensor166            Features tensor on each rank(GPU)167        optimizer: optimizer168            Optimizer for partial fc169 170        Returns:171        --------172        x_grad: tensor173            The gradient of features.174        loss_v: tensor175            Loss value for cross entropy.176        """177        total_label, norm_weight = self.prepare(label, optimizer)178        total_features = torch.zeros(179            size=[self.batch_size * self.world_size, self.embedding_size], device=self.device)180        dist.all_gather(list(total_features.chunk(self.world_size, dim=0)), features.data)181        total_features.requires_grad = True182 183        logits = self.forward(total_features, norm_weight)184        logits = self.margin_softmax(logits, total_label)185 186        with torch.no_grad():187            max_fc = torch.max(logits, dim=1, keepdim=True)[0]188            dist.all_reduce(max_fc, dist.ReduceOp.MAX)189 190            # calculate exp(logits) and all-reduce191            logits_exp = torch.exp(logits - max_fc)192            logits_sum_exp = logits_exp.sum(dim=1, keepdims=True)193            dist.all_reduce(logits_sum_exp, dist.ReduceOp.SUM)194 195            # calculate prob196            logits_exp.div_(logits_sum_exp)197 198            # get one-hot199            grad = logits_exp200            index = torch.where(total_label != -1)[0]201            one_hot = torch.zeros(size=[index.size()[0], grad.size()[1]], device=grad.device)202            one_hot.scatter_(1, total_label[index, None], 1)203 204            # calculate loss205            loss = torch.zeros(grad.size()[0], 1, device=grad.device)206            loss[index] = grad[index].gather(1, total_label[index, None])207            dist.all_reduce(loss, dist.ReduceOp.SUM)208            loss_v = loss.clamp_min_(1e-30).log_().mean() * (-1)209 210            # calculate grad211            grad[index] -= one_hot212            grad.div_(self.batch_size * self.world_size)213 214        logits.backward(grad)215        if total_features.grad is not None:216            total_features.grad.detach_()217        x_grad: torch.Tensor = torch.zeros_like(features, requires_grad=True)218        # feature gradient all-reduce219        dist.reduce_scatter(x_grad, list(total_features.grad.chunk(self.world_size, dim=0)))220        x_grad = x_grad * self.world_size221        # backward backbone222        return x_grad, loss_v223