Upload 3 files

pretrain.py

@@ -0,0 +1,270 @@
+from performer_pytorch import PerformerLM
+from performer_pytorch.autoregressive_wrapper import AutoregressiveWrapper
+
+import argparse
+import random
+import os
+from tqdm import tqdm
+import gzip
+import numpy as np
+import torch
+import torch.optim as optim
+from torch.nn import functional as F
+from torch.utils.data import DataLoader, Dataset
+from torch.cuda.amp import autocast, GradScaler
+
+from functools import reduce
+import pandas as pd
+from scipy import sparse
+from sklearn.model_selection import train_test_split, ShuffleSplit, StratifiedShuffleSplit, StratifiedKFold
+from sklearn.metrics import accuracy_score, f1_score, confusion_matrix, precision_recall_fscore_support, classification_report
+from torch import nn
+from torch.optim import Adam, SGD, AdamW
+from torch.optim.lr_scheduler import StepLR, CosineAnnealingWarmRestarts, CyclicLR
+from torch.utils.data import DataLoader, Dataset
+from torch.utils.data.distributed import DistributedSampler
+from torch.nn.parallel import DistributedDataParallel as DDP
+import torch.distributed as dist
+
+import scanpy as sc
+import anndata as ad
+from utils import *
+import pickle as pkl
+
+from sophia import SophiaG
+
+
+os.environ['CUDA_LAUNCH_BLOCKING'] = '1'
+
+# # constants
+
+# NUM_BATCHES = int(1e5)
+# BATCH_SIZE = 4
+GRADIENT_ACCUMULATE_EVERY = 4
+LEARNING_RATE = 1e-4
+VALIDATE_EVERY  = 100
+GENERATE_EVERY  = 500
+# GENERATE_LENGTH = 2048
+# SEQ_LEN = 4096
+
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--local_rank", type=int, default=-1, help='Local process rank.')
+parser.add_argument("--bin_num", type=int, default=5, help='Number of bins.')
+parser.add_argument("--gene_num", type=int, default=16906, help='Number of genes.')
+parser.add_argument("--epoch", type=int, default=1, help='Number of epochs.')
+parser.add_argument("--seed", type=int, default=2021, help='Random seed.')
+parser.add_argument("--batch_size", type=int, default=8, help='Number of batch size.')
+parser.add_argument("--learning_rate", type=float, default=1e-4, help='Learning rate.')
+parser.add_argument("--grad_acc", type=int, default=60, help='Number of gradient accumulation.')
+parser.add_argument("--valid_every", type=int, default=1, help='Number of training epochs between twice validation.')
+parser.add_argument("--pos_embed", type=bool, default=True, help='Using Gene2vec encoding or not.')
+parser.add_argument("--data_path", type=str, default='./data/panglao_human.h5ad', help='Path of data for finetune.')
+parser.add_argument("--model_path", type=str, default='./panglao_pretrained.pth', help='Path of pretrained model.')
+parser.add_argument("--ckpt_dir", type=str, default='./ckpts/', help='Directory of checkpoint to save.')
+parser.add_argument("--model_name", type=str, default='finetune', help='Finetuned model name.')
+
+args = parser.parse_args()
+# rank = int(os.environ["RANK"])
+# local_rank = args.local_rank
+# is_master = local_rank == 0
+
+SEED = args.seed
+EPOCHS = args.epoch
+BATCH_SIZE = args.batch_size
+GRADIENT_ACCUMULATION = args.grad_acc
+LEARNING_RATE = args.learning_rate
+SEQ_LEN = args.gene_num + 1
+VALIDATE_EVERY = args.valid_every
+
+PATIENCE = 10
+UNASSIGN_THRES = 0.0
+
+CLASS = args.bin_num + 2
+POS_EMBED_USING = args.pos_embed
+
+model_name = args.model_name
+ckpt_dir = args.ckpt_dir
+
+# dist.init_process_group(backend='nccl')
+# torch.cuda.set_device(local_rank)
+# device = torch.device("cuda", local_rank)
+# world_size = torch.distributed.get_world_size()
+
+# seed_all(SEED + torch.distributed.get_rank())
+
+
+
+# helpers
+
+def cycle(loader):
+    while True:
+        for data in loader:
+            yield data
+
+def decode_token(token):
+    return str(chr(max(32, token)))
+
+def decode_tokens(tokens):
+    return ''.join(list(map(decode_token, tokens)))
+
+# instantiate model
+
+model = PerformerLM(
+    num_tokens = args.bin_num + 2,
+    dim = 200,
+    depth = 3,
+    max_seq_len = SEQ_LEN,
+    heads = 5,
+    causal = False,
+    reversible = False,
+    use_scalenorm = True,
+    local_attn_heads = 0,
+    g2v_position_emb = POS_EMBED_USING,
+    generalized_attention = True
+)
+
+model = AutoregressiveWrapper(model)
+model.cuda()
+
+
+
+# prepare sc data
+
+class SCDataset(Dataset):
+    def __init__(self, data, label):
+        super().__init__()
+        self.data = data
+        self.label = label
+
+    def __getitem__(self, index):
+        rand_start = random.randint(0, self.data.shape[0]-1)
+        full_seq = self.data[rand_start].toarray()[0]
+        full_seq[full_seq > (CLASS - 2)] = CLASS - 2
+        full_seq = torch.from_numpy(full_seq).long()
+        full_seq = torch.cat((full_seq, torch.tensor([0]))).to(device)
+        seq_label = self.label[rand_start]
+        return full_seq, seq_label
+
+    def __len__(self):
+        return self.data.shape[0]
+    
+class SCDatasetPretrain(Dataset):
+    def __init__(self, data, seq_len):
+        super().__init__()
+        self.data = data
+        self.seq_len = seq_len
+
+    def __getitem__(self, index):
+        # rand_start = torch.randint(0, self.data.size(0) - self.seq_len - 1, (1,))
+        # full_seq = self.data[rand_start: rand_start + self.seq_len + 1].long()
+        
+        rand_start = random.randint(0, self.data.shape[0]-1)
+        full_seq = self.data[rand_start].toarray()[0]
+        full_seq[full_seq > (CLASS - 2)] = CLASS - 2
+        full_seq = torch.from_numpy(full_seq).long()
+        full_seq = torch.cat((full_seq, torch.tensor([0])))
+        
+        return full_seq.cuda()
+
+    def __len__(self):
+        return self.data.shape[0]
+    
+
+data = sc.read_h5ad(args.data_path)
+#data = data[:1000, :]
+# label_dict, label = np.unique(np.array(data.obs['cell_type']), return_inverse=True)  # Convert strings categorical to integrate categorical, and label_dict[label] can be restored
+# #store the label dict and label for prediction
+# with open('label_dict', 'wb') as fp:
+#     pkl.dump(label_dict, fp)
+# with open('label', 'wb') as fp:
+#     pkl.dump(label, fp)
+# class_num = np.unique(label, return_counts=True)[1].tolist()
+# class_weight = torch.tensor([(1 - (x / sum(class_num))) ** 2 for x in class_num])
+# label = torch.from_numpy(label)
+data = data.X
+
+acc = []
+f1 = []
+f1w = []
+skf = StratifiedKFold(n_splits=5, shuffle=True, random_state=SEED)
+pred_list = pd.Series(['un'] * data.shape[0])
+
+# sss = StratifiedShuffleSplit(n_splits=1, test_size=0.2, random_state=SEED)
+# for index_train in sss.split(data):
+#     data_train = data[index_train]
+#     data_val = data[index_val]
+#     train_dataset = SCDatasetPretrain(data_train, SEQ_LEN)
+#     val_dataset = SCDatasetPretrain(data_val, SEQ_LEN)
+
+# train_sampler = DistributedSampler(train_dataset)
+# val_sampler = DistributedSampler(val_dataset)
+# train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, sampler=train_sampler)
+# val_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE, sampler=val_sampler)
+
+index_train = int(data.shape[0]*0.8)
+data_train = data[:index_train]
+data_val = data[index_train:]
+train_dataset = SCDatasetPretrain(data_train, SEQ_LEN)
+val_dataset = SCDatasetPretrain(data_val, SEQ_LEN)
+
+train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE)
+val_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE)
+# train_loader  = cycle(DataLoader(train_dataset, batch_size = BATCH_SIZE))
+# val_loader    = cycle(DataLoader(val_dataset, batch_size = BATCH_SIZE))
+
+# optimizer
+
+optim = SophiaG(model.parameters(), lr=2e-4, 
+                    betas=(0.965, 0.99), rho = 0.01, weight_decay=1e-1)
+# optim = torch.optim.SGD(model.parameters(), lr=1e-8, momentum=0.9)
+# optim = torch.optim.Adam(model.parameters(), lr=LEARNING_RATE)
+scaler = GradScaler()
+
+# training
+
+for i in tqdm(range(EPOCHS), mininterval=10., desc='training'):
+    model.train()
+
+    # for __ in range(GRADIENT_ACCUMULATE_EVERY):
+    with autocast():
+        # loss = model(next(train_loader), return_loss = True)
+        for index, data_batch in enumerate(tqdm(train_loader)):
+            loss = model(data_batch, return_loss = True)
+            #print(f'training loss: {loss.item()}')
+                
+        scaler.scale(loss).backward()
+        #print(f'training loss: {loss.item()}')
+
+    print(f'training loss: {loss.item()}')
+
+    scaler.unscale_(optim)
+    torch.nn.utils.clip_grad_norm_(model.parameters(), 0.5)
+    scaler.step(optim)
+    scaler.update()
+    optim.zero_grad()
+
+    # if i % VALIDATE_EVERY == 0:
+    #     model.eval()
+    #     with torch.no_grad():
+    #         #loss = model(next(val_loader), return_loss = True)
+    #         for index, data_batch in enumerate(tqdm(val_loader)):
+    #             loss = model(data_batch, return_loss = True)
+    #             print(f'validation loss: {loss.item()}')
+
+    if i % GENERATE_EVERY == 0 and i != 0:
+        model.eval()
+        inp = random.choice(val_dataset)[:-1]
+        prime = decode_tokens(inp)
+        print(f'%s \n\n %s', (prime, '*' * 100))
+
+        sample = model.generate(inp, GENERATE_LENGTH)
+        output_str = decode_tokens(sample)
+        print(output_str)
+
+# save model
+print('save model')
+checkpoint = {'state_dict': model.state_dict(),'optimizer' :optim.state_dict()}
+torch.save(checkpoint, os.path.join(ckpt_dir, 'model_gene_attn.pth'))
+
+a=1

sophia.py

@@ -0,0 +1,202 @@
+import math
+import torch
+from torch import Tensor
+from torch.optim.optimizer import Optimizer
+from typing import List, Optional
+
+
+class SophiaG(Optimizer):
+    def __init__(self, params, lr=1e-4, betas=(0.965, 0.99), rho = 0.04,
+         weight_decay=1e-1, *, maximize: bool = False,
+         capturable: bool = False):
+        if not 0.0 <= lr:
+            raise ValueError("Invalid learning rate: {}".format(lr))
+        if not 0.0 <= betas[0] < 1.0:
+            raise ValueError("Invalid beta parameter at index 0: {}".format(betas[0]))
+        if not 0.0 <= betas[1] < 1.0:
+            raise ValueError("Invalid beta parameter at index 1: {}".format(betas[1]))
+        if not 0.0 <= rho:
+            raise ValueError("Invalid rho parameter at index 1: {}".format(rho))
+        if not 0.0 <= weight_decay:
+            raise ValueError("Invalid weight_decay value: {}".format(weight_decay))
+        defaults = dict(lr=lr, betas=betas, rho=rho, 
+                        weight_decay=weight_decay, 
+                        maximize=maximize, capturable=capturable)
+        super(SophiaG, self).__init__(params, defaults)
+
+    def __setstate__(self, state):
+        super().__setstate__(state)
+        for group in self.param_groups:
+            group.setdefault('maximize', False)
+            group.setdefault('capturable', False)
+        state_values = list(self.state.values())
+        step_is_tensor = (len(state_values) != 0) and torch.is_tensor(state_values[0]['step'])
+        if not step_is_tensor:
+            for s in state_values:
+                s['step'] = torch.tensor(float(s['step']))
+    
+    @torch.no_grad()
+    def update_hessian(self):
+        for group in self.param_groups:
+            beta1, beta2 = group['betas']
+            for p in group['params']:
+                if p.grad is None:
+                    continue
+                state = self.state[p]
+
+                if len(state) == 0:
+                    state['step'] = torch.zeros((1,), dtype=torch.float, device=p.device) \
+                        if self.defaults['capturable'] else torch.tensor(0.)
+                    state['exp_avg'] = torch.zeros_like(p, memory_format=torch.preserve_format)
+                    state['hessian'] = torch.zeros_like(p, memory_format=torch.preserve_format)
+                
+                if 'hessian' not in state.keys():
+                    state['hessian'] = torch.zeros_like(p, memory_format=torch.preserve_format)
+
+                state['hessian'].mul_(beta2).addcmul_(p.grad, p.grad, value=1 - beta2)
+
+
+    @torch.no_grad()
+    def step(self, closure=None, bs=5120):
+        loss = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+
+        for group in self.param_groups:
+            params_with_grad = []
+            grads = []
+            exp_avgs = []
+            state_steps = []
+            hessian = []
+            beta1, beta2 = group['betas']
+
+            for p in group['params']:
+                if p.grad is None:
+                    continue
+                params_with_grad.append(p)
+                
+                if p.grad.is_sparse:
+                    raise RuntimeError('Hero does not support sparse gradients')
+                grads.append(p.grad)
+                state = self.state[p]
+                # State initialization
+                if len(state) == 0:
+                    state['step'] = torch.zeros((1,), dtype=torch.float, device=p.device) \
+                        if self.defaults['capturable'] else torch.tensor(0.)
+                    state['exp_avg'] = torch.zeros_like(p, memory_format=torch.preserve_format)
+                    state['hessian'] = torch.zeros_like(p, memory_format=torch.preserve_format)
+                
+                if 'hessian' not in state.keys():
+                    state['hessian'] = torch.zeros_like(p, memory_format=torch.preserve_format)                
+
+                exp_avgs.append(state['exp_avg'])
+                state_steps.append(state['step'])
+                hessian.append(state['hessian'])
+                
+                if self.defaults['capturable']:
+                    bs = torch.ones((1,), dtype=torch.float, device=p.device) * bs
+
+            sophiag(params_with_grad,
+                  grads,
+                  exp_avgs,
+                  hessian,
+                  state_steps,
+                  bs=bs,
+                  beta1=beta1,
+                  beta2=beta2,
+                  rho=group['rho'],
+                  lr=group['lr'],
+                  weight_decay=group['weight_decay'],
+                  maximize=group['maximize'],
+                  capturable=group['capturable'])
+
+        return loss
+
+def sophiag(params: List[Tensor],
+          grads: List[Tensor],
+          exp_avgs: List[Tensor],
+          hessian: List[Tensor],
+          state_steps: List[Tensor],
+          capturable: bool = False,
+          *,
+          bs: int,
+          beta1: float,
+          beta2: float,
+          rho: float,
+          lr: float,
+          weight_decay: float,
+          maximize: bool):
+
+    if not all(isinstance(t, torch.Tensor) for t in state_steps):
+        raise RuntimeError("API has changed, `state_steps` argument must contain a list of singleton tensors")
+
+    
+    func = _single_tensor_sophiag
+
+    func(params,
+         grads,
+         exp_avgs,
+         hessian,
+         state_steps,
+         bs=bs,
+         beta1=beta1,
+         beta2=beta2,
+         rho=rho,
+         lr=lr,
+         weight_decay=weight_decay,
+         maximize=maximize,
+         capturable=capturable)
+
+def _single_tensor_sophiag(params: List[Tensor],
+                         grads: List[Tensor],
+                         exp_avgs: List[Tensor],
+                         hessian: List[Tensor],
+                         state_steps: List[Tensor],
+                         *,
+                         bs: int,
+                         beta1: float,
+                         beta2: float,
+                         rho: float,
+                         lr: float,
+                         weight_decay: float,
+                         maximize: bool,
+                         capturable: bool):
+
+    for i, param in enumerate(params):
+        grad = grads[i] if not maximize else -grads[i]
+        exp_avg = exp_avgs[i]
+        hess = hessian[i]
+        step_t = state_steps[i]
+
+        if capturable:
+            assert param.is_cuda and step_t.is_cuda and bs.is_cuda 
+            
+        if torch.is_complex(param):
+            grad = torch.view_as_real(grad)
+            exp_avg = torch.view_as_real(exp_avg)
+            hess = torch.view_as_real(hess)
+            param = torch.view_as_real(param)
+
+        # update step
+        step_t += 1
+
+        # Perform stepweight decay
+        param.mul_(1 - lr * weight_decay)
+
+        # Decay the first and second moment running average coefficient
+        exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)
+        
+        if capturable:
+            step = step_t
+            step_size = lr 
+            step_size_neg = step_size.neg()
+
+            ratio = (exp_avg.abs() / (rho * bs * hess + 1e-15)).clamp(None,1)
+            param.addcmul_(exp_avg.sign(), ratio, value=step_size_neg)
+        else:
+            step = step_t.item()
+            step_size_neg = - lr 
+            
+            ratio = (exp_avg.abs() / (rho * bs * hess + 1e-15)).clamp(None,1)
+            param.addcmul_(exp_avg.sign(), ratio, value=step_size_neg)

utils.py

@@ -0,0 +1,376 @@
+# -*- coding: utf-8 -*-
+
+from __future__ import print_function
+import json
+import os
+import struct
+import sys
+import platform
+import re
+import time
+import traceback
+import requests
+import socket
+import random
+import math
+import numpy as np
+import torch
+import logging
+import datetime
+from torch.optim.lr_scheduler import _LRScheduler
+from torch import nn
+import torch.nn.functional as F
+from torch.nn.modules.loss import _WeightedLoss
+
+
+
+def seed_all(seed_value, cuda_deterministic=False):
+    """
+    设置所有的随机种子
+    """
+    random.seed(seed_value)
+    os.environ['PYTHONHASHSEED'] = str(seed_value)
+    np.random.seed(seed_value)
+    torch.manual_seed(seed_value)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed(seed_value)
+        torch.cuda.manual_seed_all(seed_value)
+    # Speed-reproducibility tradeoff https://pytorch.org/docs/stable/notes/randomness.html
+    if cuda_deterministic:  # slower, more reproducible
+        torch.backends.cudnn.deterministic = True
+        torch.backends.cudnn.benchmark = False
+    else:  # faster, less reproducible
+        torch.backends.cudnn.deterministic = False
+        torch.backends.cudnn.benchmark = True
+
+
+def set_log(logfileName, rank=-1):
+    """
+    master节点保存所有log，其他节点只保存warning及error
+    """
+    log_file_folder = os.path.dirname(logfileName)
+    time_now = datetime.datetime.now()
+    logfileName = f'{logfileName}_{time_now.year}_{time_now.month}_{time_now.day}_{time_now.hour}_{time_now.minute}.log'
+    if not os.path.exists(log_file_folder):
+        os.makedirs(log_file_folder)
+    else:
+        pass
+
+    logging.basicConfig(level=logging.INFO if rank in [-1, 0] else logging.WARN,
+        format='[%(asctime)s %(levelname)s %(filename)s line %(lineno)d %(process)d] %(message)s',
+        datefmt='[%X]',
+        handlers=[logging.FileHandler(logfileName), logging.StreamHandler()]
+    )
+    logger = logging.getLogger()
+    return logger
+
+
+def save_ckpt(epoch, model, optimizer, scheduler, losses, model_name, ckpt_folder):
+    """
+    保存模型checkpoint
+    """
+    if not os.path.exists(ckpt_folder):
+        os.makedirs(ckpt_folder)
+    torch.save(
+        {
+            'epoch': epoch,
+            'model_state_dict': model.module.state_dict(),
+            'optimizer_state_dict': optimizer.state_dict(),
+            'scheduler_state_dict': scheduler.state_dict(),
+            'losses': losses,
+        },
+        f'{ckpt_folder}{model_name}_{epoch}.pth'
+    )
+
+def save_simple_ckpt(model, model_name, ckpt_folder):
+    """
+    保存模型checkpoint
+    """
+    if not os.path.exists(ckpt_folder):
+        os.makedirs(ckpt_folder)
+    torch.save(
+        {
+            'model_state_dict': model.module.state_dict()
+        },
+        f'{ckpt_folder}{model_name}.pth'
+    )
+
+def save_best_ckpt(epoch, model, optimizer, scheduler, losses, model_name, ckpt_folder):
+    """
+    保存模型checkpoint
+    """
+    if not os.path.exists(ckpt_folder):
+        os.makedirs(ckpt_folder)
+    torch.save(
+        {
+            'epoch': epoch,
+            'model_state_dict': model.module.state_dict(),
+            'optimizer_state_dict': optimizer.state_dict(),
+            'scheduler_state_dict': scheduler.state_dict(),
+            'losses': losses,
+        },
+        f'{ckpt_folder}{model_name}_best.pth'
+    )
+
+def get_reduced(tensor, current_device, dest_device, world_size):
+    """
+    将不同GPU上的变量或tensor集中在主GPU上，并得到均值
+    """
+    tensor = tensor.clone().detach() if torch.is_tensor(tensor) else torch.tensor(tensor)
+    tensor = tensor.to(current_device)
+    torch.distributed.reduce(tensor, dst=dest_device)
+    tensor_mean = tensor.item() / world_size
+    return tensor_mean
+
+def get_ndtensor_reduced(tensor, current_device, dest_device, world_size):
+    """
+    将不同GPU上的变量或tensor集中在主GPU上，并得到均值, 需要是2维张量
+    """
+    tensor = tensor.clone().detach() if torch.is_tensor(tensor) else torch.tensor(tensor)
+    tensor = tensor.to(current_device)
+    torch.distributed.reduce(tensor, dst=dest_device)
+    tensor_mean = torch.zeros(tensor.shape)
+    if len(tensor.shape) == 2:
+        for i in range(tensor.shape[0]):
+            for j in range(tensor.shape[1]):
+                tensor_mean[i,j] = tensor[i,j].item() / world_size
+    elif len(tensor.shape) == 1:
+        for i in range(tensor.shape[0]):
+            tensor_mean[i] = tensor[i].item() / world_size
+    return tensor_mean
+
+def numel(m: torch.nn.Module, only_trainable: bool = False):
+    """
+    returns the total number of parameters used by `m` (only counting
+    shared parameters once); if `only_trainable` is True, then only
+    includes parameters with `requires_grad = True`
+    """
+    parameters = m.parameters()
+    if only_trainable:
+        parameters = list(p for p in parameters if p.requires_grad)
+    unique = dict((p.data_ptr(), p) for p in parameters).values()
+    return sum(p.numel() for p in unique)
+
+
+def label_smooth(y, K, epsilon=0.1):
+    """
+    Label smoothing for multiclass labels
+    One hot encode labels `y` over `K` classes. `y` should be of the form [1, 6, 3, etc.]
+    """
+    m = len(y)
+    out = np.ones((m, K)) * epsilon / K
+    for index in range(m):
+        out[index][y[index] - 1] += 1 - epsilon
+    return torch.tensor(out)
+
+
+class SequentialDistributedSampler(torch.utils.data.sampler.Sampler):
+    """
+    Distributed Sampler that subsamples indicies sequentially,
+    making it easier to collate all results at the end.
+    Even though we only use this sampler for eval and predict (no training),
+    which means that the model params won't have to be synced (i.e. will not hang
+    for synchronization even if varied number of forward passes), we still add extra
+    samples to the sampler to make it evenly divisible (like in `DistributedSampler`)
+    to make it easy to `gather` or `reduce` resulting tensors at the end of the loop.
+    """
+
+    def __init__(self, dataset, batch_size, world_size, rank=None, num_replicas=None):
+        if num_replicas is None:
+            if not torch.distributed.is_available():
+                raise RuntimeError("Requires distributed package to be available")
+            num_replicas = world_size
+        if rank is None:
+            if not torch.distributed.is_available():
+                raise RuntimeError("Requires distributed package to be available")
+            rank = torch.distributed.get_rank()
+        self.dataset = dataset
+        self.num_replicas = num_replicas
+        self.rank = rank
+        self.batch_size = batch_size
+        self.num_samples = int(math.ceil(len(self.dataset) * 1.0 / self.batch_size / self.num_replicas)) * self.batch_size
+        self.total_size = self.num_samples * self.num_replicas
+
+    def __iter__(self):
+        indices = list(range(len(self.dataset)))
+        # add extra samples to make it evenly divisible
+        indices += [indices[-1]] * (self.total_size - len(indices))
+        # subsample
+        indices = indices[self.rank * self.num_samples : (self.rank + 1) * self.num_samples]
+        return iter(indices)
+
+    def __len__(self):
+        return self.num_samples
+
+
+def distributed_concat(tensor, num_total_examples, world_size):
+    """
+    合并不同进程的inference结果
+    """
+    output_tensors = [tensor.clone() for _ in range(world_size)]
+    torch.distributed.all_gather(output_tensors, tensor)
+    concat = torch.cat(output_tensors, dim=0)
+    # truncate the dummy elements added by SequentialDistributedSampler
+    return concat[:num_total_examples]
+
+
+class CosineAnnealingWarmupRestarts(_LRScheduler):
+    """
+        optimizer (Optimizer): Wrapped optimizer.
+        first_cycle_steps (int): First cycle step size.
+        cycle_mult(float): Cycle steps magnification. Default: -1.
+        max_lr(float): First cycle's max learning rate. Default: 0.1.
+        min_lr(float): Min learning rate. Default: 0.001.
+        warmup_steps(int): Linear warmup step size. Default: 0.
+        gamma(float): Decrease rate of max learning rate by cycle. Default: 1.
+        last_epoch (int): The index of last epoch. Default: -1.
+    """
+    
+    def __init__(self,
+                 optimizer : torch.optim.Optimizer,
+                 first_cycle_steps : int,
+                 cycle_mult : float = 1.,
+                 max_lr : float = 0.1,
+                 min_lr : float = 0.001,
+                 warmup_steps : int = 0,
+                 gamma : float = 1.,
+                 last_epoch : int = -1
+        ):
+        assert warmup_steps < first_cycle_steps
+        
+        self.first_cycle_steps = first_cycle_steps # first cycle step size
+        self.cycle_mult = cycle_mult # cycle steps magnification
+        self.base_max_lr = max_lr # first max learning rate
+        self.max_lr = max_lr # max learning rate in the current cycle
+        self.min_lr = min_lr # min learning rate
+        self.warmup_steps = warmup_steps # warmup step size
+        self.gamma = gamma # decrease rate of max learning rate by cycle
+        
+        self.cur_cycle_steps = first_cycle_steps # first cycle step size
+        self.cycle = 0 # cycle count
+        self.step_in_cycle = last_epoch # step size of the current cycle
+        
+        super(CosineAnnealingWarmupRestarts, self).__init__(optimizer, last_epoch)
+        
+        # set learning rate min_lr
+        self.init_lr()
+
+    def init_lr(self):
+        self.base_lrs = []
+        for param_group in self.optimizer.param_groups:
+            param_group['lr'] = self.min_lr
+            self.base_lrs.append(self.min_lr)
+    
+    def get_lr(self):
+        if self.step_in_cycle == -1:
+            return self.base_lrs
+        elif self.step_in_cycle < self.warmup_steps:
+            return [(self.max_lr - base_lr)*self.step_in_cycle / self.warmup_steps + base_lr for base_lr in self.base_lrs]
+        else:
+            return [base_lr + (self.max_lr - base_lr) \
+                    * (1 + math.cos(math.pi * (self.step_in_cycle-self.warmup_steps) \
+                                    / (self.cur_cycle_steps - self.warmup_steps))) / 2
+                    for base_lr in self.base_lrs]
+
+    def step(self, epoch=None):
+        if epoch is None:
+            epoch = self.last_epoch + 1
+            self.step_in_cycle = self.step_in_cycle + 1
+            if self.step_in_cycle >= self.cur_cycle_steps:
+                self.cycle += 1
+                self.step_in_cycle = self.step_in_cycle - self.cur_cycle_steps
+                self.cur_cycle_steps = int((self.cur_cycle_steps - self.warmup_steps) * self.cycle_mult) + self.warmup_steps
+        else:
+            if epoch >= self.first_cycle_steps:
+                if self.cycle_mult == 1.:
+                    self.step_in_cycle = epoch % self.first_cycle_steps
+                    self.cycle = epoch // self.first_cycle_steps
+                else:
+                    n = int(math.log((epoch / self.first_cycle_steps * (self.cycle_mult - 1) + 1), self.cycle_mult))
+                    self.cycle = n
+                    self.step_in_cycle = epoch - int(self.first_cycle_steps * (self.cycle_mult ** n - 1) / (self.cycle_mult - 1))
+                    self.cur_cycle_steps = self.first_cycle_steps * self.cycle_mult ** (n)
+            else:
+                self.cur_cycle_steps = self.first_cycle_steps
+                self.step_in_cycle = epoch
+                
+        self.max_lr = self.base_max_lr * (self.gamma**self.cycle)
+        self.last_epoch = math.floor(epoch)
+        for param_group, lr in zip(self.optimizer.param_groups, self.get_lr()):
+            param_group['lr'] = lr
+
+
+class DistanceLoss(_WeightedLoss):
+    """
+    CrossEntropyLoss with Distance Weighted
+    """
+    def __init__(self, weight=None, reduction='mean', ignore_index = None):
+        super().__init__(weight=weight, reduction=reduction)
+        self.weight = weight
+        self.reduction = reduction
+        self.ignore_index = ignore_index
+    def forward(self, inputs, targets):
+        if len(inputs.shape) > 2:
+            inputs = inputs.reshape(-1, inputs.size(-1))
+        if len(targets.shape) > 1:
+            targets = targets.reshape(-1)
+        if self.ignore_index is not None:
+            keep_index = (targets != self.ignore_index).nonzero(as_tuple=True)[0]
+            targets = torch.index_select(targets, 0, keep_index) #targets[targets != self.ignore_index]
+            inputs = torch.index_select(inputs, 0, keep_index)
+        lsm = F.log_softmax(inputs, -1)
+        targets = torch.empty(size=(targets.size(0), inputs.size(-1)), device=targets.device).fill_(0).scatter_(1, targets.data.unsqueeze(1), 1)
+        if self.weight is not None:
+            lsm = lsm * self.weight.unsqueeze(0)
+        loss = -(targets * lsm).sum(-1)
+        inputs = nn.Softmax(dim=-1)(inputs)[..., 1:-1].argmax(dim=-1) + 1
+        # print('inputs', inputs.device, inputs.shape)
+        targets = nn.Softmax(dim=-1)(targets)[..., 1:-1].argmax(dim=-1) + 1
+        # print('targets', targets.device, targets.shape)
+        distance = abs(inputs - targets) + 1e-2
+        # print('loss.shape', loss.shape)
+        # print('distance.shape', distance.shape)
+        loss = loss * distance
+        if self.reduction == 'sum':
+            loss = loss.sum()
+        elif self.reduction == 'mean':
+            loss = loss.mean()
+        return loss
+
+
+class LabelSmoothCrossEntropyLoss(_WeightedLoss):
+    """
+    CrossEntropyLoss with Label Somoothing
+    """
+    def __init__(self, weight=None, reduction='mean', smoothing=0.0):
+        super().__init__(weight=weight, reduction=reduction)
+        self.smoothing = smoothing
+        self.weight = weight
+        self.reduction = reduction
+
+    @staticmethod
+    def _smooth_one_hot(targets: torch.Tensor, n_classes: int, smoothing=0.0):
+        assert 0 <= smoothing < 1
+        with torch.no_grad():
+            targets = torch.empty(size=(targets.size(0), n_classes),
+                                  device=targets.device) \
+                .fill_(smoothing / (n_classes - 1)) \
+                .scatter_(1, targets.data.unsqueeze(1), 1. - smoothing)
+        return targets
+
+    def forward(self, inputs, targets):
+        targets = LabelSmoothCrossEntropyLoss._smooth_one_hot(targets, inputs.size(-1),
+                                                              self.smoothing)
+        lsm = F.log_softmax(inputs, -1)
+
+        if self.weight is not None:
+            lsm = lsm * self.weight.unsqueeze(0)
+
+        loss = -(targets * lsm).sum(-1)
+
+        if self.reduction == 'sum':
+            loss = loss.sum()
+        elif self.reduction == 'mean':
+            loss = loss.mean()
+
+        return loss