如何加速计算？

当模型很大时，如何加速计算（一）

pytorch多GPU并行训练

一般来说，多GPU训练方法包括模型的并行（如AlexNet），数据并行（复制模型，数据不同，扩大了batch_size）。当然，随着GPU 的增加，训练速度的提升也是递减的，毕竟多GPU 之间通信也会占用一部分资源。
在每个GPU训练step结束后，将每块GPU的损失梯度求平均。如果不用同步BN，而是每个设备计算自己的批次数据的均值方差，效果与单GPU一致，仅仅能提升训练速度；如果使用同步BN，效果会有一定提升，但是会损失一部分并行速度。

DistributedDataParallel

# 初始化进程环境
def init_distributed_mode(args):
    # 如果是多机多卡的机器，WORLD_SIZE代表使用的机器数，RANK对应第几台机器
    # 如果是单机多卡的机器，WORLD_SIZE代表有几块GPU，RANK和LOCAL_RANK代表第几块GPU
    if'RANK'in os.environ and'WORLD_SIZE'in os.environ:
        args.rank = int(os.environ["RANK"])
        args.world_size = int(os.environ['WORLD_SIZE'])
        # LOCAL_RANK代表某个机器上第几块GPU
        args.gpu = int(os.environ['LOCAL_RANK'])
    elif'SLURM_PROCID'in os.environ:
        args.rank = int(os.environ['SLURM_PROCID'])
        args.gpu = args.rank % torch.cuda.device_count()
    else:
        print('Not using distributed mode')
        args.distributed = False
        return

    args.distributed = True

    torch.cuda.set_device(args.gpu)  # 对当前进程指定使用的GPU
    args.dist_backend = 'nccl'# 通信后端，nvidia GPU推荐使用NCCL
    dist.barrier()  # 等待每个GPU都运行完这个地方以后再继续

# 学习率倍增
def main(args):
	    if torch.cuda.is_available() isFalse:
	        raise EnvironmentError("not find GPU device for training.")
	
	    # 初始化各进程环境
	    init_distributed_mode(args=args)
	
	    rank = args.rank
	    device = torch.device(args.device)
	    batch_size = args.batch_size
	    num_classes = args.num_classes
	    weights_path = args.weights
	    args.lr *= args.world_size  # 学习率要根据并行GPU的数倍增

然后进行数据取样，在多个GPU分配数据集时首先打乱顺序，然后将数据均匀分部到各个卡上：

#给每个rank对应的进程分配训练的样本索引
train_sampler=torch.utils.data.distributed.DistributedSampler(train_data_set)
val_sampler=torch.utils.data.distributed.DistributedSampler(val_data_set)
#将样本索引每batch_size个元素组成一个list
train_batch_sampler=torch.utils.data.BatchSampler(
train_sampler,batch_size,drop_last=True)

 train_loader = torch.utils.data.DataLoader(train_data_set,
                                               batch_sampler=train_batch_sampler,
                                               pin_memory=True,   # 直接加载到显存中，达到加速效果
                                               num_workers=nw,
                                               collate_fn=train_data_set.collate_fn)

 val_loader = torch.utils.data.DataLoader(val_data_set,
                                             batch_size=batch_size,
                                             sampler=val_sampler,
                                             pin_memory=True,
                                             num_workers=nw,
                                             collate_fn=val_data_set.collate_fn)

如果有预训练权重的话，需要保证每块GPU加载的权重是一模一样的。需要在主进程保存模型初始化权重，在不同设备上载入主进程保存的权重。这样才能保证每块GOU上加载的权重是一致的：

# 实例化模型
    model = resnet34(num_classes=num_classes).to(device)

    # 如果存在预训练权重则载入
    if os.path.exists(weights_path):
        weights_dict = torch.load(weights_path, map_location=device)
        # 简单对比每层的权重参数个数是否一致
        load_weights_dict = {k: v for k, v in weights_dict.items()
                             if model.state_dict()[k].numel() == v.numel()}
        model.load_state_dict(load_weights_dict, strict=False)
    else:
        checkpoint_path = os.path.join(tempfile.gettempdir(), "initial_weights.pt")
        # 如果不存在预训练权重，需要将第一个进程中的权重保存，然后其他进程载入，保持初始化权重一致
        if rank == 0:
            torch.save(model.state_dict(), checkpoint_path)

        dist.barrier()
        # 这里注意，一定要指定map_location参数，否则会导致第一块GPU占用更多资源
        model.load_state_dict(torch.load(checkpoint_path, map_location=device))

是否冻结权重&同步BN&DDP模型：

# 是否冻结权重
    if args.freeze_layers:
        for name, para in model.named_parameters():
            # 除最后的全连接层外，其他权重全部冻结
            if"fc"notin name:
                para.requires_grad_(False)
    else:
        # 只有训练带有BN结构的网络时使用SyncBatchNorm采用意义
        if args.syncBN:
            # 使用SyncBatchNorm后训练会更耗时
            model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model).to(device)
	
	# 转为DDP模型
         model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu])

	 # optimizer使用SGD+余弦淬火策略
	    pg = [p for p in model.parameters() if p.requires_grad]
	    optimizer = optim.SGD(pg, lr=args.lr, momentum=0.9, weight_decay=0.005)
	    lf = lambda x: ((1 + math.cos(x * math.pi / args.epochs)) / 2) * (1 - args.lrf) + args.lrf  # cosine
    scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf)
 

与单GPU不同的地方：rain_sampler.set_epoch(epoch)，这行代码会在每次迭代的时候获得一个不同的生成器，每一轮开始迭代获取数据之前设置随机种子，通过改变传进的epoch参数改变打乱数据顺序。通过设置不同的随机种子，可以让不同GPU每轮拿到的数据不同。后面的部分和单GPU相同。

for epoch in range(args.epochs):
	        train_sampler.set_epoch(epoch)  

	        mean_loss = train_one_epoch(model=model,
	                                    optimizer=optimizer,
	                                    data_loader=train_loader,
	                                    device=device,
	                                    epoch=epoch)
	
	        scheduler.step()
	
	        sum_num = evaluate(model=model,
	                           data_loader=val_loader,
	                           device=device)
	        acc = sum_num / val_sampler.total_size
def train_one_epoch(model, optimizer, data_loader, device, epoch):
	    model.train()
	    loss_function = torch.nn.CrossEntropyLoss()
	    mean_loss = torch.zeros(1).to(device)
	    optimizer.zero_grad()
	
	    # 在进程0中打印训练进度
	    if is_main_process():
	        data_loader = tqdm(data_loader)
	
	    for step, data in enumerate(data_loader):
	        images, labels = data
	
	        pred = model(images.to(device))
	
	        loss = loss_function(pred, labels.to(device))
	        loss.backward()
	        loss = reduce_value(loss, average=True)  #  在单GPU中不起作用，多GPU时，获得所有GPU的loss的均值。
	        mean_loss = (mean_loss * step + loss.detach()) / (step + 1)  # update mean losses
	
	        # 在进程0中打印平均loss
	        if is_main_process():
	            data_loader.desc = "[epoch {}] mean loss {}".format(epoch, round(mean_loss.item(), 3))
	
	        ifnot torch.isfinite(loss):
	            print('WARNING: non-finite loss, ending training ', loss)
	            sys.exit(1)
	
	        optimizer.step()
	        optimizer.zero_grad()
	
	    # 等待所有进程计算完毕
	    if device != torch.device("cpu"):
	        torch.cuda.synchronize(device)
	
	    return mean_loss.item()
	
	def reduce_value(value, average=True):
	    world_size = get_world_size()
	    if world_size < 2:  # 单GPU的情况
	        return value
	
	    with torch.no_grad():
	        dist.all_reduce(value)   # 对不同设备之间的value求和
	        if average:  # 如果需要求平均，获得多块GPU计算loss的均值
	            value /= world_size
	
        return value

验证阶段与模型保存：

	@torch.no_grad()
	def evaluate(model, data_loader, device):
	    model.eval()
	
	    # 用于存储预测正确的样本个数，每块GPU都会计算自己正确样本的数量
	    sum_num = torch.zeros(1).to(device)
	
	    # 在进程0中打印验证进度
	    if is_main_process():
	        data_loader = tqdm(data_loader)
	
	    for step, data in enumerate(data_loader):
	        images, labels = data
	        pred = model(images.to(device))
	        pred = torch.max(pred, dim=1)[1]
	        sum_num += torch.eq(pred, labels.to(device)).sum()
	
	    # 等待所有进程计算完毕
	    if device != torch.device("cpu"):
	        torch.cuda.synchronize(device)
	
	    sum_num = reduce_value(sum_num, average=False)  # 预测正确样本个数
	
    return sum_num.item()
 # 模型需要在主进程中进行保存
 if rank == 0:
            print("[epoch {}] accuracy: {}".format(epoch, round(acc, 3)))
            tags = ["loss", "accuracy", "learning_rate"]
            tb_writer.add_scalar(tags[0], mean_loss, epoch)
            tb_writer.add_scalar(tags[1], acc, epoch)
            tb_writer.add_scalar(tags[2], optimizer.param_groups[0]["lr"], epoch)

            torch.save(model.module.state_dict(), "./weights/model-{}.pth".format(epoch))

如果从头开始训练，主进程生成的初始化权重是以临时文件的形式保存，需要训练完后移除掉。最后还需要撤销进程组。

if rank == 0:# 删除临时缓存文件        if os.path.exists(checkpoint_path) is True:            os.remove(checkpoint_path)    dist.destroy_process_group()  # 撤销进程组，释放资源

内容参考

python常见并行库

DASK

dask是python的一个并行计算库，可以动态调度资源提供并行计算，并行化的数据集成提供接口给numpy，pandas或者python迭代器，Task Graph 任务图非常清晰，使得开发人员和用户都可以自由地构建复杂的算法，并处理大多数数据工程框架中常见的map/filter/groupby范式难以处理的混乱情况。
图片地址：官方手册

dask示例

我们测试一下在numpy和dask中数据计算的时间。

import time
import numpy as np
import matplotlib.pyplot as plt
import dask
import dask.array as da
data = np.random.normal(loc =10.0 , scale= 0.1,size = (20000,20000))

生成一波固定期望方差的数组data，使用numpy在数据上按某个维度求该维度向量上的均值。

def numpy_compute():
    print(data.mean(axis=0))

同时使用dask对数组进行分块，这里我们每块为1000*1000。

data_p=da.from_array(data,chunks=(1000,1000))

计算dask数组中按某个维度求向量均值。

def dask_compute():
    print(data_p.mean(axis=0).compute()) 

比较二者时间：

def run_time(fun):

    start_time = time.time()
    fun()
    end_time = time.time()

    print("程序运行时间为：{} 秒".format(str(round((end_time - start_time), 1))))
    return end_time - start_time
@run_time
def task():
    numpy_compute()
@run_time
def task():
    dask_compute()

课件dask相比之下运行时间要快一些。

dask的功能远不知与此，还有许多可以并行化的操作。

multiprocessing

multiprocessing是python编写多进程的程序，支持子进程、通信和共享数据、执行不同形式的同步，提供了Process、Queue、Pipe、Lock等组件。

# 参数
multiprocessing.Process(group=None, target=None, name=None, args=(), kwargs={})
- group：分组，实际上很少使用
- target：表示调用对象，你可以传入方法的名字
- name：别名，相当于给这个进程取一个名字
- args：表示被调用对象的位置参数元组，比如target是函数a，他有两个参数m，n，那么args就传入(m, n)即可
- kwargs：表示调用对象的字典

简单使用

import math
import datetime
import multiprocessing as mp
def train_on_parameter(name, param):
    result = 0
    for num in param:
        result += math.sqrt(num * math.tanh(num) / math.log2(num) / math.log10(num))
    return {name: result}
if __name__ == '__main__':
    start_t = datetime.datetime.now()
    num_cores = int(mp.cpu_count())
    print("本地计算机有: " + str(num_cores) + " 核心")
    pool = mp.Pool(num_cores)
    param_dict = {'task1': list(range(10, 30000000)),
                  'task2': list(range(30000000, 60000000)),
                  'task3': list(range(60000000, 90000000)),
                  'task4': list(range(90000000, 120000000)),
                  'task5': list(range(120000000, 150000000)),
                  'task6': list(range(150000000, 180000000)),
                  'task7': list(range(180000000, 210000000)),
                  'task8': list(range(210000000, 240000000))}
    results = [pool.apply_async(train_on_parameter, args=(name, param)) for name, param in param_dict.items()]
    results = [p.get() for p in results]
    end_t = datetime.datetime.now()
    elapsed_sec = (end_t - start_t).total_seconds()
    print("多进程计算 共消耗: " + "{:.2f}".format(elapsed_sec) + " 秒")

运行结果：

线程池

import random
from multiprocessing.pool import Pool
from time import sleep, time
import os
def run(name):
    print("%s子进程开始，进程ID：%d" % (name, os.getpid()))
    start = time()
    sleep(random.choice([1, 2, 3, 4]))
    end = time()
    print("%s子进程结束，进程ID：%d。耗时0.2%f" % (name, os.getpid(), end-start))
if __name__ == "__main__":
    print("父进程开始")
    # 创建多个进程，表示可以同时执行的进程数量。默认大小是CPU的核心数
    p = Pool(8)
    for i in range(10):
        # 创建进程，放入进程池统一管理
        p.apply_async(run, args=(i,))
    # 如果我们用的是进程池，在调用join()之前必须要先close()，并且在close()之后不能再继续往进程池添加新的进程
    p.close()
    # 进程池对象调用join，会等待进程吃中所有的子进程结束完毕再去结束父进程
    p.join()
    print("父进程结束。")

threading

Thread 对象数据属性有name（线程名）,ident（线程标识符），daemon（线程是否是守护线程）等。主要对象包括start（），run()和join（）等。start表示开始执行该线程，run()定义线程功能，通常在子类中被应用开发者重写，join (timeout=None)表示直到启动的线程终止之前一直挂起，除非给出timeout秒，否则会一直阻塞。

import threading
import time
def read():
    for x in range(5):
        print('在%s,正在听音乐' % time.ctime())
        time.sleep(1.5)
def write():
    for x in range(5):
        print('在%s,正在看电视' % time.ctime())
        time.sleep(1.5)

def main():
    music_threads = []  # 用来存放执行read函数线程的列表
    TV_threads = []  # 用来存放执行write函数线程的列表
    for i in range(1,2):  # 创建1个线程用于read()，并添加到read_threads列表
        t = threading.Thread(target=read) # 执行的函数如果需要传递参数，threading.Thread(target=函数名,args=(参数，逗号隔开))
        music_threads.append(t)
    for i in range(1,2): # 创建1个线程执行write()，并添加到write_threads列表
        t = threading.Thread(target=write) # 执行的函数如果需要传递参数，threading.Thread(target=函数名,args=(参数，逗号隔开))
        TV_threads.append(t)
    for i in range(0,1):  # 启动存放在read_threads和write_threads列表中的线程
        music_threads[i].start()
        TV_threads[i].start()
if __name__ == '__main__':
    main()

本例参考博客

为了让线程更好的封装，可以用threading模块下的Thread类，继承这个类，然后实现run方法，线程就回自动运行run方法中的代码。

import threading
import time
count = 0
class MyThread(threading.Thread):
    def __init__(self , threadName):
        super(MyThread,self).__init__(name=threadName)
    """一旦这个MyThread类被调用，自动的就会运行底下的run方法中的代码，
    因为这个run方法所属的的MyThread类继承了threading.Thread"""
    def run(self):
        global count
        for i in range(100):
            count += 1
            time.sleep(0.3)
            print(self.getName() , count)
for i in range(2):
    MyThread("MyThreadName:" + str(i)).start()