PyTorch的GPU训练方式学习

GPU并行训练

模型转cuda

gpus=[0]#指定训练的GPU
cuda_gpu=torch.cuda.is_available()   #判断GPU是否存在可用

net=Net()#模型初始化
if(cuda_gpu):
	net=torch.nn.DataParallel(net,device_ids=gps).cuda()#ids指定GPU

数据转cuda

(minibatchX, minibatchY) = minibatch
minibatchX = minibatchX.astype(np.float32).T#转浮点型数字
minibatchY = minibatchY.astype(np.float32).T
if(cuda_gpu):
    b_x = Variable(torch.from_numpy(minibatchX).cuda())    #将数据转为cuda类型
    b_y = Variable(torch.from_numpy(minibatchY).cuda())
else:
    b_x = Variable(torch.from_numpy(minibatchX))
    b_y = Variable(torch.from_numpy(minibatchY))

输出数据去cuda，转为numpy

correct_prediction = sum(torch.max(output, 1)[1].data.squeeze() == torch.max(b_y, 1)[1].data.squeeze())
if(cuda_gpu):
    correct_prediction = correct_prediction.cpu().numpy()   #.cpu将cuda转为tensor类型，.numpy将tensor转为numpy类型
else:
    correct_prediction = correct_prediction.numpy()

单机多卡 分布式训练

补充文章：DDP入门教程
DDP实现原理与源代码解析
DDP实战技巧

torch.distributed 包支持

分布训练包括 GPU 和 CPU 的分布式训练支持。Pytorch 分布式目前只支持 Linux。
torch.distributed 分布式训练的优势：

1.每个进程对应一个独立的训练过程，且只对梯度等少量数据进行信息交换。

• 在每次迭代中，每个进程具有自己的 optimizer ，并独立完成所有的优化步骤，进程内与一般的训练无异。
• 在各进程梯度计算完成之后，各进程需要将梯度进行汇总平均，然后再由 rank=0 的进程，将其 broadcast(广播) 到所有进程。之后，各进程用该梯度来更新参数
• 由于各进程中的模型，初始参数一致 (初始时刻进行一次 broadcast)，而每次用于更新参数的梯度也一致，因此，各进程的模型参数始终保持一致
  而在 DataParallel 中，全程维护一个 optimizer，对各 GPU 上梯度进行求和，而在主 GPU 进行参数更新，之后再将模型参数 broadcast 到其他 GPU。所以分布式训练的参数传输的数据量更少，因此速度更快效率更高。
  基本概念
• group
  即进程组。默认情况下，只有一个组，一个 job 即为一个组，也即一个 world。当需要进行更加精细的通信时，可以通过 new_group 接口，使用 word 的子集，创建新组，用于集体通信等。
• world size
  全局进程个数
• rank
  进程序号，用于进程间通讯，表征进程优先级。rank=0的主机为mask
• local_rank
  进程内GPU编号，非显式参数，由 torch.distributed.launch 内部指定。比方说， rank = 3，local_rank = 0 表示第 3 个进程内的第 1 块 GPU。
  基本使用流程

1. 在使用 distributed 包的任何其他函数之前，需要使用 init_process_group 初始化进程组，同时初始化 distributed 包

distributed.init_process_group(backend='nccl',rank,world_size,init_method='env://')
#Env初始化方式

2. 如果需要进行小组内集体通信，用 new_group 创建子分组
3. 创建分布式并行模型 DDP(model, device_ids=device_ids)
4. 为数据集创建 Sampler

train_sample=DistributeedSampler(traindataset)

5. 使用启动工具 torch.distributed.launch --n在每个主机上执行一次脚本，开始训练
6. 使用 destory_process_group() 销毁进程组

编程模板
一,Env初始化方法

import torch.distributed as dist
import torch.utils.data.distributed

# ......
import argparse
parser = argparse.ArgumentParser()
# 注意这个参数，必须要以这种形式指定，即使代码中不使用。因为 launch 工具默认传递该参数
parser.add_argument("--local_rank", type=int)
args = parser.parse_args()

# ......
dist.init_process_group(backend='nccl', init_method='env://')

# ......
trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=download, transform=transform)
train_sampler = torch.utils.data.distributed.DistributedSampler(trainset)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size, sampler=train_sampler)

# ......
# 根据 local_rank，配置当前进程使用的 GPU
net = Net()
device = torch.device('cuda', args.local_rank)
net = net.to(device)
net = torch.nn.parallel.DistributedDataParallel(net, device_ids=[args.local_rank], output_device=args.local_rank)

执行方式

python -m torch.distributed.launch --nproc_per_node=2 --nnodes=3 --node_rank=0 --master_addr="192.168.1.201" --master_port=23456 env_init.py

python -m torch.distributed.launch --nproc_per_node=2 --nnodes=3 --node_rank=1 --master_addr="192.168.1.201" --master_port=23456 env_init.py

python -m torch.distributed.launch --nproc_per_node=2 --nnodes=3 --node_rank=2 --master_addr="192.168.1.201" --master_port=23456 env_init.py

其中，节点数为3，每个节点的进程数为2，node_rank是每个节点的优先级，–master_addr是指定主节点

说明

1. 在 Env 方式中，在 init_process_group 中，无需指定任何参数
2. 必须在rank==0的进程内保存参数
3. 该方式下，使用 torch.distributed.launch 在每台主机上，为其创建多进程，其中:
4. nproc_per_node 参数指定为当前主机创建的进程数。一般设定为当前主机的 GPU 数量
5. nnodes 参数指定当前 job 包含多少个节点
6. node_rank 指定当前节点的优先级
7. master_addr 和 master_port 分别指定 master 节点的 ip:port
8. 若没有为每个进程合理分配 GPU，则默认使用当前主机上所有的 GPU。即使一台主机上有多个进程，也会共用 GPU
9. 使用 torch.distributed.launch 工具时，将会为当前主机创建 nproc_per_node 个进程，每个进程独立执行训练脚本。同时，它还会为每个进程分配一个 local_rank 参数，表示当前进程在当前主机上的编号。例如：rank=2, local_rank=0 表示第 3 个节点上的第 1 个进程。
10. 需要合理利用 local_rank 参数，来合理分配本地的 GPU 资源
11. 每条命令表示一个进程。若已开启的进程未达到 word_size 的数量，则所有进程会一直等待
    二,TCP初始化方式
    代码：

import torch.distributed as dist
import torch.utils.data.distributed

# ......
parser = argparse.ArgumentParser(description='PyTorch distributed training on cifar-10')
parser.add_argument('--rank', default=0,
                    help='rank of current process')
parser.add_argument('--word_size', default=2,
                    help="word size")
parser.add_argument('--init_method', default='tcp://127.0.0.1:23456',
                    help="init-method")
args = parser.parse_args()

# ......
dist.init_process_group(backend='nccl', init_method=args.init_method, rank=args.rank, world_size=args.word_size)

# ......
trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=download, transform=transform)
train_sampler = torch.utils.data.distributed.DistributedSampler(trainset)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size, sampler=train_sampler)

# ......
net = Net()
net = net.cuda()
net = torch.nn.parallel.DistributedDataParallel(net)

执行方式：

# Node 1 : ip 192.168.1.201  port : 12345
python tcp_init.py --init_method tcp://192.168.1.201:12345 --rank 0 --word_size 3

# Node 2 : 
python tcp_init.py --init_method tcp://192.168.1.201:12345 --rank 1 --word_size 3

# Node 3 : 
python tcp_init.py --init_method tcp://192.168.1.201:12345 --rank 2 --word_size 3

说明：

1. 在 TCP 方式中，在 init_process_group 中必须手动指定以下参数
2. rank 为当前进程的进程号
3. word_size 为当前 job 的总进程数
4. init_method 内指定 tcp 模式，且所有进程的 ip:port 必须一致，设定为主进程的 ip:port
5. 必须在 rank==0 的进程内保存参数
6. 若程序内未根据 rank 设定当前进程使用的 GPUs，则默认使用全部 GPU，且以数据并行的方式使用。
7. 每条命令表示一个进程。若已开启的进程未达到 word_size 的数量，则所有进程会一直等待
8. 每台主机上可以开启多个进程。但是，若未为每个进程分配合适的 GPU，则同机不同进程可能会共用 GPU，应该坚决避免这种情况。
9. 使用 gloo 后端进行 GPU 训练时，会报错。
10. 若每个进程负责多块 GPU，可以利用多 GPU 进行模型并行。如下所示：

class ToyMpModel(nn.Module):
    def init(self, dev0, dev1):
        super(ToyMpModel, self).init()
        self.dev0 = dev0
        self.dev1 = dev1
        self.net1 = torch.nn.Linear(10, 10).to(dev0)
        self.relu = torch.nn.ReLU()
        self.net2 = torch.nn.Linear(10, 5).to(dev1)

def forward(self, x):
       x = x.to(self.dev0)
       x = self.relu(self.net1(x))
       x = x.to(self.dev1)
       return self.net2(x)

......
dev0 = rank * 2
dev1 = rank * 2 + 1
mp_model = ToyMpModel(dev0, dev1)#模型并行
ddp_mp_model = DDP(mp_model)

进程组说明
1.初始化进程组

torch.distributed.init_process_group(backend,init_method=None, 
                                     timeout=datetime.timedelta(0, 1800), 
                                     world_size=-1, 
                                     rank=-1, 
                                     store=None)

2,参数说明

• backend:指定当前进程要使用的通信后端
  小写字符串，支持的通信后端有 gloo，mpi，nccl 。建议用 nccl。具体含义后面说明

• init_method ： 指定当前进程组初始化方式

• rank: 指定当前进程的优先级

• world_size:该 job 中的总进程数。如果指定 store 参数，则需要指定该参数。

• timeout:指定每个进程的超时时间
  可选参数，datetime.timedelta 对象，默认为 30 分钟。该参数仅用于 Gloo 后端

• store
  所有 worker 可访问的 key / value，用于交换连接 / 地址信息。与 init_method 互斥。
  new_group

torch.distributed.new_group(ranks=None, 
                            timeout=datetime.timedelta(0, 1800), 
                            backend=None)

• 函数作用
  new_group() 函数可用于使用所有进程的任意子集来创建新组。其返回一个分组句柄，可作为 collectives 相关函数的 group 参数 。collectives 是分布式函数，用于特定编程模式中的信息交换。
• 参数详解
  ranks:指定新分组内的成员的 ranks 列表
  timeout：指定该分组进程组内的操作的超时时间
  backend：指定要使用的通信后端,小写字符串，支持的通信后端有 gloo，nccl ，必须与 init_process_group() 中一致。
  获取进程组属性
• get_backend
• get_rank
• get_world_size
• is_initialized
• is_mpi_available
• is_nccl_available

通信后端

• Gloo后端
  gloo 后端支持 CPU 和 GPU，其支持集体通信（collective Communication），并对其进行了优化。 由于 GPU 之间可以直接进行数据交换，而无需经过 CPU 和内存，因此，在 GPU 上使用 gloo 后端速度更快。torch.distributed 对 gloo 提供原生支持，无需进行额外操作。
• NCCL后端
  NCCL 的全称为 Nvidia 聚合通信库（NVIDIA Collective Communications Library），是一个可以实现多个 GPU、多个结点间聚合通信的库，在 PCIe、Nvlink、InfiniBand 上可以实现较高的通信速度。NCCL 高度优化和兼容了 MPI，并且可以感知 GPU 的拓扑，促进多 GPU 多节点的加速，最大化 GPU 内的带宽利用率，所以深度学习框架的研究员可以利用 NCCL 的这个优势，在多个结点内或者跨界点间可以充分利用所有可利用的 GPU。NCCL 对 CPU 和 GPU 均有较好支持，且 torch.distributed 对其也提供了原生支持。对于每台主机均使用多进程的情况，使用 NCCL 可以获得最大化的性能。每个进程内，不许对其使用的 GPUs 具有独占权。若进程之间共享 GPUs 资源，则可能导致 deadlocks。
• MPI后端
  MPI 即消息传递接口（Message Passing Interface），是一个来自于高性能计算领域的标准的工具。它支持点对点通信以及集体通信，并且是 torch.distributed 的 API 的灵感来源。使用 MPI 后端的优势在于，在大型计算机集群上，MPI 应用广泛，且高度优化。但是，torch.distributed 对 MPI 并不提供原生支持。因此，要使用 MPI，必须从源码编译 Pytorch。是否支持 GPU，视安装的 MPI 版本而定。
  后端如何选择
  基本原则：
  用 NCCL 进行分布式 GPU 训练
  用 Gloo 进行分布式 CPU 训练
  更加具体的教程，参考：https://zhuanlan.zhihu.com/p/76638962
  https://www.jiqizhixin.com/articles/2019-04-30-8
  https://zhuanlan.zhihu.com/p/113694038

示例代码，多GPU的训练代码

# -*- coding=utf-8 -*-

import pdb

import argparse
import os
import os.path as osp
import random
import shutil
import time
import warnings
import sys
import torch
import torch.nn as nn
import torch.nn.parallel
import torch.nn.functional as F
import torch.backends.cudnn as cudnn

# import torch.optim

# from torch.optim.lr_scheduler import MultiStepLR

from torch.utils import data
import torch.utils.data.distributed
import torchvision.transforms as transforms
import torchvision.datasets as datasets
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from collections import OrderedDict
from easydict import EasyDict as edict
# 添加当前工程内的package
cur_dir = osp.split(osp.realpath(__file__))[0]
src_dir = osp.join(cur_dir, '../')
if src_dir not in sys.path:
    sys.path.append(src_dir)
from src.dataset.sampler import DistributedTestSampler
from src.dataset.datasets_3d import video_data
from src.model.inceptionv2_3D_len8 import inceptionv2_3D
from src.model.criterion import center_loss
from src.utils.log import Log
from src.dataset import group_transform
from src.utils.utils import save_checkpoint, AverageMeter, accuracy,data_prefetcher, group_weight_3d
from src.utils.acc import acc_com, acc_thr
from src.utils.model_load import model_load_p3d, model_load_len8,
model_load, model_load_test
import re
C = edict()
config = 
#
=======================================================================================================================
# common config
exp_time = time.strftime('%Y_%m_%d_%H_%M_%S', time.localtime())
C.file_name = osp.splitext(osp.split(__file__)[1])[0]

C.snapshot_dir =r'/data1/video_classification/2d/scripts/model_3d_gv2_base/'
# train and test file config
C.mean = [0.485, 0.456, 0.406]
C.std = [0.229, 0.224, 0.225]
C.mean_cv = [123,117,104]#[104, 117, 123]
C.std_cv = [0.017, 0.017, 0.017]
C.test_file='/data1/video_classification/2d/list/video_cls_test_list_20190520_3c.txt'
C.train_file='/data1/video_classification/2d/list/train_list_len8_20190716_new_3c.txt'
C.train_img_path = r'/dataset/zhangruixuan3/'
C.test_img_path = r'/dataset/zhangruixuan3/'
# Setting for Network
C.pretrained = True
# C.pretrained_model_inceptionv2 = r'../src/pre_model/inception_maxpooling_len8_iter_80000.pth'
C.pretrained_model_inceptionv2=r'/data1/model_analysis/video/4L_32F_Att5/26pth/checkpoint_E26.pth'
# Train Config
C.lr = 0.001
C.momentum = 0.9
C.weight_decay = 0.00005
C.batch_size = 64  # 12 # all gpu
C.num_workers = 5
C.start_epoch = 1
C.epochs = 40
# C.mile_stones = [20,40]
C.snapshot_epoch = list(range(1, C.epochs)) + [C.epochs]  # [5, 10, 15, 20, 25, 30]
print(C.snapshot_epoch)
C.print_freq = 100
def main():
    parser = argparse.ArgumentParser()
    cudnn.benchmark = True
    parser.add_argument('--local_rank', default=0, type=int, help='processrank on node')
    parser.add_argument('--center_loss', default=1, type=int, help='processrank on node')
    parser.add_argument('--weight_center_loss',default=0.5, type=float, help='center_loss')
    args = parser.parse_args()
    seed = args.local_rank
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed(seed)


    # 初始化distributed environment
	torch.cuda.set_device(args.local_rank)
    os.environ['MASTER_ADDR'] ='10.3.67.12'
    os.environ['MASTER_PORT'] ='58397'
    dist.init_process_group(backend='nccl', rank=0, world_size=1,init_method='env://')
    if args.local_rank == 0:
        os.makedirs(C.snapshot_dir,exist_ok=True)
        log =Log(osp.join(C.snapshot_dir, 'log_DDP.txt'))
        args.log = log
 

    args.file_name = C.file_name
    args.ckpt_dir = C.snapshot_dir
    args.world_size =dist.get_world_size()
    #####inceptionv3
    train_transform =transforms.Compose([
       group_transform.GroupScale((224, 224), cv=True),#cv=False),
       group_transform.GroupRandomHorizontalFlip(),# randomly reverse from leftto right
       group_transform.Stack(),
       group_transform.ToTorchFormatTensor(div=False),#True
       group_transform.GroupNormalize(mean=C.mean_cv, std=C.std_cv, cv=True)

    ])

    test_transform =transforms.Compose([
        # transforms.Resize((256,256)),
        # transforms.CenterCrop((224, 224)),
       group_transform.GroupScale((224, 224), cv=True),group_transform.Stack(),
       group_transform.ToTorchFormatTensor(div=False),
       group_transform.GroupNormalize(mean=C.mean_cv, std=C.std_cv, cv=True)

    ])

 
    train_set = video_data(
        img_path=C.train_img_path,
        txt_path=C.train_file,
        num_segments=1,
        length=8,
        modality='RGB',
        img_tmpl='image_{:04d}.jpg',
        step=2,
        length_first=True,
        data_transforms=train_transform)



    train_sampler=torch.utils.data.distributed.DistributedSampler(train_set)

    train_loader =torch.utils.data.DataLoader(
        dataset=train_set,
        batch_size=C.batch_size,
        shuffle=False,
        num_workers=C.num_workers,
        pin_memory=True,
        drop_last=False,
        sampler=train_sampler)


    # test

    test_set = video_data(
        img_path=C.test_img_path,
        txt_path=C.test_file,
        num_segments=1,
        length=8,
        modality='RGB',
        img_tmpl='image_{:04d}.jpg',
        step=2,
        length_first=True,
        data_transforms=test_transform)

    test_sampler =DistributedTestSampler(test_set)
    test_loader = torch.utils.data.DataLoader(
        dataset=test_set,
        batch_size=C.batch_size,
        shuffle=False,
        num_workers=C.num_workers,
        pin_memory=True,
        drop_last=False,
        sampler=test_sampler)
    model = inceptionv2_3D(num_classes=3,use_attention=True)


    if C.pretrained:

        # model =model_load_len8(C.pretrained_model_inceptionv2, model)

        model =model_load_test(C.pretrained_model_inceptionv2, model)
    # move to corresponding GPU
   	model.cuda()
    model = DDP(model,device_ids=[args.local_rank]
    print('rank: {}, model scatteringdone'.format(args.local_rank))
    criterion_CE =nn.CrossEntropyLoss(ignore_index=-1).cuda(device=args.local_rank)
    weight_group = []
    group_weight_3d(weight_group,module=model, norm_layer=None, lr=C.lr)

    ##

    total_params_id = list(map(id,model.module.parameters()))
    fc_id = list(map(id,model.module.fc_new.parameters()))
    base_params_id = [x for x in total_params_id if x not in fc_id]
    total_params =list(model.module.parameters())
    fc_params = [x for x intotal_params if id(x) in fc_id]
    base_params = [x for x intotal_params if id(x) in base_params_id]
    optimizer = torch.optim.SGD([
        {'params': base_params, 'lr':C.lr},
        {'params': fc_params, 'lr':C.lr * 10}],
        momentum=C.momentum,weight_decay=C.weight_decay)

    # set center loss

    #if args.center_loss:

    #    center_criterion =center_loss.CenterLoss(num_classes=2, feat_dim=1024, use_gpu=True,
    #                              pretrained_path=C.pretrained_model_inceptionv2)
    #    optimizer_center_loss =torch.optim.SGD(center_criterion.parameters(), lr=C.lr)

 

    # pdb.set_trace()

    # init_lr=[x['lr'] for x in optimizer.state_dict()['param_groups']]

    #per_iter_num=len(train_loader.batch_sampler)

    # # 对不同参数设置不同的学习率
    # optimizer = torch.optim.SGD(
    #     model.module.parameters(),lr=0.1 ,
    #     momentum = C.momentum, weight_decay=C.weight_decay)
    # 设置学习率调整策略
    # scheduler =MultiStepLR(optimizer, milestones=mile_stones)
    scheduler =torch.optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.1)
    # 开始训练
    for epoch in range(C.start_epoch,C.epochs + 1):

       train_sampler.set_epoch(epoch)

        # train for one epoch

       train_with_prefeatch(train_loader, model, criterion_CE, optimizer,epoch, scheduler, args)

 

        # # evaluate on validation set

        validate_DDP(test_loader,model, epoch, args)
        if (args.local_rank == 0) and(epoch in C.snapshot_epoch):
            args.log('saving checkpoint...')
            save_checkpoint({
                'epoch': epoch,
                'arch':'InceptionV2',
                'state_dict':model.state_dict(),
                'optimizer':optimizer.state_dict(),
                #'center_dict':center_criterion._parameters,

            }, False,filepath=C.snapshot_dir)



def train_with_prefeatch(train_loader, model, criterion, optimizer, epoch,scheduler, args):
    cudnn.beckmark = True
    batch_time = AverageMeter()
    data_time = AverageMeter()
    losses = AverageMeter()

    #center_losses = AverageMeter()
    top1 = AverageMeter()
    top5 = AverageMeter(
    # 按需调整学习率
    scheduler.step()
    #
max_iter_num=per_iter_num*C.epochs

    # switch to train mode
    model.train()
    end = time.time()
    epoch_start = time.time()

    prefetcher =data_prefetcher(train_loader, device=args.local_rank)
    input, target = prefetcher.next()

    if args.local_rank == 0:
        print(len(target),target[0].size())
    i = -1
    while input is not None:

        i += 1

        # print (i)

        data_time.update(time.time()- end)


        output, feat =model(input[0])

        assert (len(output) ==len(target))

        loss =torch.zeros([1]).cuda(args.local_rank)
        prec1_mean =torch.zeros([1]).cuda(args.local_rank)
        loss_CE = criterion(output,target)
        loss += loss_CE
        # loss_center,loss_others_to_centers, loss_center_to_center = center_criterion(feat, target)

        # print('loss', loss_center,loss_center_to_center, loss_others_to_centers)

        # loss += loss_CE +args.weight_center_loss * (loss_center + loss_others_to_centers +loss_center_to_center)

        acc1, acc5 = accuracy(output,target, topk=(1, 2))

        losses.update(loss.item(),input[0].size(0))
      
#center_losses.update(loss_center.item(), input[0].size(0))

        #others_to_centers_losses.update(loss_others_to_centers.item(),input[0].size(0))


#center_to_center_losses.update(loss_center_to_center.item(),input[0].size(0))

        # prec1_mean, _ =accuracy(output[0].data, target.data, topk=(1, 2))

        reduced_loss =reduce_tensor(loss.data)
        acc1 = reduce_tensor(acc1)

        acc5 = reduce_tensor(acc5)

        #reduced_center_loss =reduce_tensor(loss_center.data)

       
#reduced_others_to_center_loss =reduce_tensor(loss_others_to_centers.data)

        #reduced_center_to_center_loss =reduce_tensor(loss_center_to_center.data)

 

       losses.update(reduced_loss.data.item(), input[0].size(0))

        top1.update(acc1[0],input[0].size(0))

        top5.update(acc5[0],input[0].size(0))

        #center_losses.update(reduced_center_loss.data.item(),input[0].size(0))

       
#others_to_centers_losses.update(reduced_others_to_center_loss.data.item(),input[0].size(0))

       
#center_to_center_losses.update(reduced_center_to_center_loss.item(), input[0].size(0))

 

        # optimizer =poly_lr_scheduler( optimizer, (epoch-1)*per_iter_num+i, max_iter_num,init_lr,power=1)

        # compute gradient and do SGD step

        optimizer.zero_grad()

        #optimizer_center_loss.zero_grad()

        loss.backward()
        optimizer.step()
        torch.cuda.synchronize()
        #for param incenter_criterion.parameters():
        #    param.grad.data *= (1. /args.weight_center_loss)
        #optimizer_center_loss.step()
        # measure elapsed time
        batch_time.update(time.time()- end)
        end = time.time()
        # featch next batch
        input, target =prefetcher.next()
        if (args.local_rank == 0) and(i % C.print_freq == 0):
            args.log('Rank: {0}\t'
            'Epoch:[{1}][{2}/{3}]\t'
            'Time{batch_time.val:.3f} ({batch_time.avg:.3f})\t'
            'Data{data_time.val:.3f} ({data_time.avg:.3f})\t'
            'Loss{loss.val:.4f} ({loss.avg:.4f})\t'
            'Acc@1{top1.val:.3f} ({top1.avg:.3f})\t'
            'Acc@5{top5.val:.3f} ({top5.avg:.3f})\t'.format(args.local_rank,epoch, i, len(train_loader), batch_time=batch_time,data_time=data_time,
loss=losses, top1=top1, top5=top5),['{:.7f}'.format(x['lr']) for x in optimizer.state_dict()['param_groups']]
    epoch_end = time.time()
    if args.local_rank == 0:
        args.log('epoch time cost:{:.2f}s'.format(epoch_end - epoch_start))
def validate_DDP(val_loader, model, epoch, args):
    # switch to evaluate mode
    model.eval()
    end = time.time()
    cudnn.beckmark = False
    probs =torch.FloatTensor().cuda(args.local_rank) 
# 需要人工做gather
    label =torch.LongTensor().cuda(args.local_rank)
    if args.local_rank == 0:
        args.log('testing {} with epoch {}'.format(osp.split(C.test_file)[-1], epoch))
    with torch.no_grad():
        for i, (input, target) in enumerate(val_loader):
            if (args.local_rank == 0)and (i % C.print_freq == 0):
                args.log('testing {}th sample, time cost: {:.2f}s'.format(i * C.batch_size, time.time() - end))
            input =input[0].cuda(device=args.local_rank, non_blocking=True)
            if isinstance(target,list) or isinstance(target, tuple):
                target =[x.cuda(device=args.local_rank, non_blocking=True) for x in target]
                target =[x.unsqueeze(dim=1) for x in target]
                target = torch.cat(target, dim=1)

            elif isinstance(target,torch.Tensor):

                target =target.cuda(device=args.local_rank, non_blocking=True)
            utput, feat =model(input)
            # 先转变成probability
            output = F.softmax(output, dim=1)
            # save current proborresponding target
            probs = torch.cat((probs,output), dim=0)
            label = torch.cat((label,target), dim=0)
            # gather probs and target
    rt_probs_list =[torch.zeros_like(probs) for i in range(dist.get_world_size())]
    rt_label_list =[torch.zeros_like(label) for i in range(dist.get_world_size())]
    dist.all_gather(rt_probs_list,probs)
    dist.all_gather(rt_label_list,label)
    probs_total = torch.cat(rt_probs_list, dim=0)
    label_total =torch.cat(rt_label_list, dim=0)
    probs_total =probs_total.data.cpu()
    label_total =label_total.data.cpu()
    if args.local_rank == 0:
       args.log('beforeprobs_total.shape: ', probs_total.shape)
       args.log('beforelabel_total.shape: ', label_total.shape)
    # 要去掉一些重复的，在DistributedTestSampler中导致的重复
    sample_num =len(val_loader.dataset)
    probs_total =probs_total[0:sample_num, :]
    label_total = label_total[0:sample_num]
    if args.local_rank == 0:
        args.log('afterprobs_total.shape: ', probs_total.shape)
        args.log('afterlabel_total.shape: ', label_total.shape)
        acc_com(probs_total,label_total, args)
def reduce_tensor(tensor, gpu_num=8):
    gpu_num = dist.get_world_size()
    rt = tensor.clone()
    dist.all_reduce(rt,
op=dist.ReduceOp.SUM)
    rt /= gpu_num
    return rt
if __name__ == '__main__':
    main()