vpl.py

import logging
import os

import torch
import torch.distributed as dist
from torch.nn import Module
from torch.nn.functional import normalize, linear
from torch.nn.parameter import Parameter


class VPL(Module):
    """
    Modified from Partial-FC
    """

    @torch.no_grad()
    def __init__(self, rank, local_rank, world_size, batch_size, resume,
                 margin_softmax, num_classes, sample_rate=1.0, embedding_size=512, prefix="./", cfg=None):
        super(VPL, self).__init__()
        #
        assert sample_rate==1.0
        assert not resume
        self.num_classes: int = num_classes
        self.rank: int = rank
        self.local_rank: int = local_rank
        self.device: torch.device = torch.device("cuda:{}".format(self.local_rank))
        self.world_size: int = world_size
        self.batch_size: int = batch_size
        self.margin_softmax: callable = margin_softmax
        self.sample_rate: float = sample_rate
        self.embedding_size: int = embedding_size
        self.prefix: str = prefix
        self.num_local: int = num_classes // world_size + int(rank < num_classes % world_size)
        self.class_start: int = num_classes // world_size * rank + min(rank, num_classes % world_size)
        self.num_sample: int = int(self.sample_rate * self.num_local)

        self.weight_name = os.path.join(self.prefix, "rank_{}_softmax_weight.pt".format(self.rank))
        self.weight_mom_name = os.path.join(self.prefix, "rank_{}_softmax_weight_mom.pt".format(self.rank))

        self.weight = torch.normal(0, 0.01, (self.num_local, self.embedding_size), device=self.device)
        self.weight_mom: torch.Tensor = torch.zeros_like(self.weight)
        logging.info("softmax weight init successfully!")
        logging.info("softmax weight mom init successfully!")
        self.stream: torch.cuda.Stream = torch.cuda.Stream(local_rank)

        self.index = None
        self.update = lambda: 0
        self.sub_weight = Parameter(self.weight)
        self.sub_weight_mom = self.weight_mom

        #vpl variables

        self._iters = 0
        self.cfg = cfg
        self.vpl_mode = -1
        if self.cfg is not None:
            self.vpl_mode = self.cfg['mode']
            if self.vpl_mode>=0:
                self.register_buffer("queue", torch.randn(self.num_local, self.embedding_size, device=self.device))
                self.queue = normalize(self.queue)
                self.register_buffer("queue_iters", torch.zeros((self.num_local,), dtype=torch.long, device=self.device))
                self.register_buffer("queue_lambda", torch.zeros((self.num_local,), dtype=torch.float32, device=self.device))


    def save_params(self):
        pass
        #torch.save(self.weight.data, self.weight_name)
        #torch.save(self.weight_mom, self.weight_mom_name)

    @torch.no_grad()
    def sample(self, total_label):
        index_positive = (self.class_start <= total_label) & (total_label < self.class_start + self.num_local)
        total_label[~index_positive] = -1
        total_label[index_positive] -= self.class_start
        return index_positive

    def forward(self, total_features, norm_weight):
        torch.cuda.current_stream().wait_stream(self.stream)
        logits = linear(total_features, norm_weight)
        return logits

    @torch.no_grad()
    def update(self):
        self.weight_mom[self.index] = self.sub_weight_mom
        self.weight[self.index] = self.sub_weight

    def prepare(self, label, optimizer):
        with torch.cuda.stream(self.stream):
            total_label = torch.zeros(
                size=[self.batch_size * self.world_size], device=self.device, dtype=torch.long)
            dist.all_gather(list(total_label.chunk(self.world_size, dim=0)), label)
            index_positive = self.sample(total_label)
            optimizer.state.pop(optimizer.param_groups[-1]['params'][0], None)
            optimizer.param_groups[-1]['params'][0] = self.sub_weight
            optimizer.state[self.sub_weight]['momentum_buffer'] = self.sub_weight_mom
            norm_weight = normalize(self.sub_weight)
            return total_label, norm_weight, index_positive

    @torch.no_grad()
    def prepare_queue_lambda(self, label, iters):
        self.queue_lambda[:] = 0.0
        if iters>self.cfg['start_iters']:
            allowed_delta = self.cfg['allowed_delta']
            if self.vpl_mode==0:
                past_iters = iters - self.queue_iters
                idx = torch.where(past_iters <= allowed_delta)[0]
                self.queue_lambda[idx] = self.cfg['lambda']

            if iters % 2000 == 0 and self.rank == 0:
                logging.info('[%d]use-lambda: %d/%d'%(iters,len(idx), self.num_local))

    @torch.no_grad()
    def set_queue(self, total_features, total_label, index_positive, iters):
        local_label = total_label[index_positive]
        sel_features = normalize(total_features[index_positive,:])
        self.queue[local_label,:] = sel_features
        self.queue_iters[local_label] = iters

    def forward_backward(self, label, features, optimizer, feature_w):
        self._iters += 1
        total_label, norm_weight, index_positive = self.prepare(label, optimizer)
        total_features = torch.zeros(
            size=[self.batch_size * self.world_size, self.embedding_size], device=self.device)
        dist.all_gather(list(total_features.chunk(self.world_size, dim=0)), features.data)
        total_features.requires_grad = True

        if feature_w is not None:
            total_feature_w = torch.zeros(
                size=[self.batch_size * self.world_size, self.embedding_size], device=self.device)
            dist.all_gather(list(total_feature_w.chunk(self.world_size, dim=0)), feature_w.data)

        if self.vpl_mode>=0:
            self.prepare_queue_lambda(total_label, self._iters)
            _lambda = self.queue_lambda.view(self.num_local, 1)
            injected_weight = norm_weight*(1.0-_lambda) + self.queue*_lambda
            injected_norm_weight = normalize(injected_weight)
            logits = self.forward(total_features, injected_norm_weight)
        else:
            logits = self.forward(total_features, norm_weight)

        logits = self.margin_softmax(logits, total_label)

        with torch.no_grad():
            max_fc = torch.max(logits, dim=1, keepdim=True)[0]
            dist.all_reduce(max_fc, dist.ReduceOp.MAX)

            # calculate exp(logits) and all-reduce
            logits_exp = torch.exp(logits - max_fc)
            logits_sum_exp = logits_exp.sum(dim=1, keepdims=True)
            dist.all_reduce(logits_sum_exp, dist.ReduceOp.SUM)

            # calculate prob
            logits_exp.div_(logits_sum_exp)

            # get one-hot
            grad = logits_exp
            index = torch.where(total_label != -1)[0]
            one_hot = torch.zeros(size=[index.size()[0], grad.size()[1]], device=grad.device)
            one_hot.scatter_(1, total_label[index, None], 1)

            # calculate loss
            loss = torch.zeros(grad.size()[0], 1, device=grad.device)
            loss[index] = grad[index].gather(1, total_label[index, None])
            dist.all_reduce(loss, dist.ReduceOp.SUM)
            loss_v = loss.clamp_min_(1e-30).log_().mean() * (-1)

            # calculate grad
            grad[index] -= one_hot
            grad.div_(self.batch_size * self.world_size)

        logits.backward(grad)
        if total_features.grad is not None:
            total_features.grad.detach_()
        x_grad: torch.Tensor = torch.zeros_like(features, requires_grad=True)
        # feature gradient all-reduce
        dist.reduce_scatter(x_grad, list(total_features.grad.chunk(self.world_size, dim=0)))
        x_grad = x_grad * self.world_size
        #vpl set queue
        if self.vpl_mode>=0:
            if feature_w is None:
                self.set_queue(total_features.detach(), total_label, index_positive, self._iters)
            else:
                self.set_queue(total_feature_w, total_label, index_positive, self._iters)
        # backward backbone
        return x_grad, loss_v