PyTorch从入门到实践 | (5) PyTorch常用工具

在训练神经网络过程中，需要用到很多工具，其中最重要的三部分是：数据加载预处理、可视化和GPU加速。本章主要介绍Pytorch在这几方面的工具模块，合理使用这些工具能够极大地提高编码效率。

目录

1. 数据处理

2. 计算机视觉工具包：torchvision

3. 可视化工具

4. 使用GPU加速：cuda

5. 持久化

---

1. 数据处理

在解决深度学习问题的过程中，往往需要花费大量的精力去处理数据，包括图像、文本、语音或其它二进制数据等。数据的处理对训练神经网络来说十分重要，良好的数据处理不仅会加速模型训练，更会提高模型效果。考虑到这点，PyTorch提供了几个高效便捷的工具，以便使用者进行数据处理或增强等操作，同时可通过并行化加速数据加载。

• 数据加载

在PyTorch中，数据加载可通过自定义的数据集对象。数据集对象被抽象为Dataset类，实现自定义的数据集需要继承Dataset，并实现两个Python魔法方法：

1）__getitem__：返回一条数据，或一个样本。obj[index]等价于obj.__getitem__(index)

2）__len__：返回样本的数量。len(obj)等价于obj.__len__()

这里我们以Kaggle经典挑战赛"Dogs vs. Cat"的数据为例，来详细讲解如何处理数据。"Dogs vs. Cats"是一个分类问题，判断一张图片是狗还是猫，其所有图片都存放在一个文件夹下，根据文件名的前缀判断是狗还是猫。

%env LS_COLORS = None 
!tree --charset ascii  data/dogcat/ #查看./data/dogcat目录结构

import torch
from torch.utils import data

import os
from PIL import  Image
import numpy as np

class DogCat(data.Dataset): #自定义数据集继承Dataset
    def __init__(self, root):
        imgs = os.listdir(root) #root下的所有文件
        # 所有图片的绝对路径
        # 这里不实际加载图片，只是指定路径，当调用__getitem__时才会真正读图片
        self.imgs = [os.path.join(root, img) for img in imgs]
      
    #需要实现以下两个魔法方法
    def __getitem__(self, index): #使用obj[index]就会调用这个方法 
        img_path = self.imgs[index]
        # dog->1， cat->0
        label = 1 if 'dog' in img_path.split('/')[-1] else 0
        pil_img = Image.open(img_path)#图像读取到数组中
        array = np.asarray(pil_img) 
        data = torch.from_numpy(array) #数组转为tensor
        return data, label
    
    def __len__(self): #使用len(obj)就会调用这个方法
        return len(self.imgs)

dataset = DogCat('./data/dogcat/')
img, label = dataset[0] # 相当于调用dataset.__getitem__(0)
for img, label in dataset:
    print(img.size(), img.float().mean(), label)

通过上面的代码，我们学习了如何自定义自己的数据集，并可以依次获取。但这里返回的数据不适合实际使用，因其具有如下两方面问题：

1）返回样本的形状不一，因每张图片的大小不一样，这对于需要取batch训练的神经网络来说很不友好

2）返回样本的数值较大，未归一化至[-1, 1]

针对上述问题，PyTorch提供了torchvisionhttps://github.com/pytorch/vision/。它是一个视觉工具包，提供了很多视觉图像处理的工具，其中transforms模块提供了对PIL Image对象和Tensor对象的常用操作。

对PIL Image的操作包括：

1）Scale：调整图片尺寸，长宽比保持不变

2）CenterCrop、RandomCrop、RandomResizedCrop： 裁剪图片

3）Pad：填充

4）ToTensor：将PIL Image对象转成Tensor，会自动将[0, 255]归一化至[0, 1]

对Tensor的操作包括：

1）Normalize：标准化，即减均值，除以标准差

2）ToPILImage：将Tensor转为PIL Image对象

如果要对图片进行多个操作，可通过Compose函数将这些操作拼接起来，类似于nn.Sequential。注意，这些操作定义后是以函数的形式存在，真正使用时需调用它的__call__方法，这点类似于nn.Module。例如要将图片调整为224×224，首先应构建这个操作trans = Resize((224, 224))，然后调用trans(img)。下面我们就用transforms的这些操作来优化上面实现的dataset：

import os
from PIL import  Image
import numpy as np
from torchvision import transforms as T

transform = T.Compose([
    T.Resize(224), # 缩放图片(Image)，保持长宽比不变，最短边为224像素
    T.CenterCrop(224), # 从图片中间切出224*224的图片
    T.ToTensor(), # 将图片(Image)转成Tensor，归一化至[0, 1] 数组->tensor 如(224,224,3) -> (3,224,224)
    T.Normalize(mean=[.5, .5, .5], std=[.5, .5, .5]) # 标准化至[-1, 1]，规定每个通道的均值和标准差
])

class DogCat(data.Dataset):
    def __init__(self, root, transforms=None):
        imgs = os.listdir(root)
        self.imgs = [os.path.join(root, img) for img in imgs]
        self.transforms=transforms
        
    def __getitem__(self, index):
        img_path = self.imgs[index]
        label = 0 if 'dog' in img_path.split('/')[-1] else 1
        data = Image.open(img_path)
        if self.transforms:
            data = self.transforms(data)
        return data, label
    
    def __len__(self):
        return len(self.imgs)

dataset = DogCat('./data/dogcat/', transforms=transform)
img, label = dataset[0]
for img, label in dataset:
    print(img.size(), label)

除了上述操作之外，transforms还可通过Lambda封装自定义的转换策略。例如想对PIL Image进行随机旋转，则可写成这样trans=T.Lambda(lambda img: img.rotate(random()*360))。

torchvision已经预先实现了常用的Dataset，包括前面使用过的CIFAR-10，以及ImageNet、COCO、MNIST、LSUN等数据集，可通过诸如torchvision.datasets.CIFAR10来调用，具体使用方法请参看官方文档。在这里介绍一个会经常使用到的Dataset——ImageFolder，它的实现和上述的DogCat很相似。ImageFolder假设所有的文件按文件夹保存，每个文件夹下存储同一个类别的图片，文件夹名为类名，其构造函数如下：

ImageFolder(root, transform=None, target_transform=None, loader=default_loader)

它主要有四个参数:

1) root：在root指定的路径下寻找图片

2) transform：对PIL Image进行的转换操作，transform的输入是使用loader读取图片的返回对象

3) target_transform：对label的转换

4) loader：给定路径后如何读取图片，默认读取为RGB格式的PIL Image对象

label是按照文件夹名顺序排序后存成字典，即{类名:类序号(从0开始)}，一般来说最好直接将文件夹命名为从0开始的数字，这样会和ImageFolder实际的label一致，如果不是这种命名规范，建议看看self.class_to_idx属性以了解label和文件夹名的映射关系。

!tree --charset ASCII  data/dogcat_2/

from torchvision.datasets import ImageFolder
dataset = ImageFolder('data/dogcat_2/') #对于上述特定格式的数据集 直接使用ImageFolder快速构建dataset

# cat文件夹的图片对应label 0，dog对应1
print(dataset.class_to_idx)
# 所有图片的路径和对应的label
print(dataset.imgs)

# 没有任何的transform，所以返回的还是PIL Image对象
print(dataset[0][1]) # 第一维指定第几张图，第二维为1返回label
dataset[0][0] # 为0返回图片数据

# 加上transform
normalize = T.Normalize(mean=[0.4, 0.4, 0.4], std=[0.2, 0.2, 0.2])
transform  = T.Compose([
         T.RandomResizedCrop(224),
         T.RandomHorizontalFlip(),
         T.ToTensor(),
         normalize,
])

dataset = ImageFolder('data/dogcat_2/', transform=transform)
# 深度学习中图片数据一般保存成CxHxW，即通道数x图片高x图片宽
print(dataset[0][0].size())
to_img = T.ToPILImage() #定义从tensor转回图像数据的操作
# 0.2和0.4是标准差和均值的近似
to_img(dataset[0][0]*0.2+0.4) #还原 标准化操作 各通道乘以标准差加上均值

Dataset只负责数据的抽象，一次调用__getitem__只返回一个样本。前面提到过，在训练神经网络时，最好是对一个batch的数据进行操作，同时还需要对数据进行shuffle和并行加速等。对此，PyTorch提供了DataLoader帮助我们实现这些功能。

DataLoader的函数定义如下：

DataLoader(dataset, batch_size=1, shuffle=False, sampler=None, num_workers=0, collate_fn=default_collate, pin_memory=False, drop_last=False)

from torch.utils.data import DataLoader

dataloader = DataLoader(dataset, batch_size=3, shuffle=True, num_workers=0, drop_last=False)

dataiter = iter(dataloader)
imgs, labels = next(dataiter) #取出一个batch
imgs.size() # batch_size, channel, height, weight

dataloader是一个可迭代的对象，意味着我们可以像使用迭代器一样使用它，例如：

在数据处理中，有时会出现某个样本无法读取等问题，比如某张图片损坏。这时在__getitem__函数中将出现异常，此时最好的解决方案即是将出错的样本剔除。如果实在是遇到这种情况无法处理，则可以返回None对象，然后在Dataloader中实现自定义的collate_fn，将空对象过滤掉。但要注意，在这种情况下dataloader返回的batch中的数目会少于batch_size。

class NewDogCat(DogCat): # 继承前面实现的DogCat数据集
    def __getitem__(self, index):
        try:
            # 调用父类的获取函数，即 DogCat.__getitem__(self, index)
            return super(NewDogCat,self).__getitem__(index)
        except:
            return None, None

from torch.utils.data.dataloader import default_collate # 导入默认的拼接方式
def my_collate_fn(batch):
    '''
    batch中每个元素形如(data, label)
    '''
    # 过滤为None的数据
    batch = list(filter(lambda x:x[0] is not None, batch))
    if len(batch) == 0: return torch.Tensor()
    return default_collate(batch) # 用默认方式拼接过滤后的batch数据

dataset = NewDogCat('data/dogcat_wrong/', transforms=transform)
dataset[5]

dataloader = DataLoader(dataset, 2, collate_fn=my_collate_fn, num_workers=1,shuffle=True)
for batch_datas, batch_labels in dataloader:
    print(batch_datas.size(),batch_labels.size())

来看一下上述batch_size的大小。其中第2个的batch_size为1，这是因为有一张图片损坏，导致其无法正常返回。而最后1个的batch_size也为1，这是因为共有9张（包括损坏的文件）图片，无法整除2（batch_size），因此最后一个batch的数据会少于batch_szie，可通过指定drop_last=True来丢弃最后一个不足batch_size的batch。

对于诸如样本损坏或数据集加载异常等情况，还可以通过其它方式解决。例如但凡遇到异常情况，就随机取一张图片代替：

相比较丢弃异常图片而言，这种做法会更好一些，因为它能保证每个batch的数目仍是batch_size。但在大多数情况下，最好的方式还是对数据进行彻底清洗。

DataLoader里面并没有太多的魔法方法，它封装了Python的标准库multiprocessing，使其能够实现多进程加速。在此提几点关于Dataset和DataLoader使用方面的建议：

1）高负载的操作放在__getitem__中，如加载图片等。

2）dataset中应尽量只包含只读对象，避免修改任何可变对象，利用多线程进行操作。

第一点是因为多进程会并行的调用__getitem__函数，将负载高的放在__getitem__函数中能够实现并行加速。 第二点是因为dataloader使用多进程加载，如果在Dataset实现中使用了可变对象，可能会有意想不到的冲突。在多线程/多进程中，修改一个可变对象，需要加锁，但是dataloader的设计使得其很难加锁（在实际使用中也应尽量避免锁的存在），因此最好避免在dataset中修改可变对象。例如下面就是一个不好的例子，在多进程处理中self.num可能与预期不符，这种问题不会报错，因此难以发现。如果一定要修改可变对象，建议使用Python标准库Queue中的相关数据结构。

使用Python multiprocessing库的另一个问题是，在使用多进程时，如果主程序异常终止（比如用Ctrl+C强行退出），相应的数据加载进程可能无法正常退出。这时你可能会发现程序已经退出了，但GPU显存和内存依旧被占用着，或通过top、ps aux依旧能够看到已经退出的程序，这时就需要手动强行杀掉进程。建议使用如下命令：

PyTorch中还单独提供了一个sampler模块，用来对数据进行采样。常用的有随机采样器：RandomSampler，当dataloader的shuffle参数为True时，系统会自动调用这个采样器，实现打乱数据。默认的是采用SequentialSampler，它会按顺序一个一个进行采样。这里介绍另外一个很有用的采样方法： WeightedRandomSampler，它会根据每个样本的权重选取数据，在样本比例不均衡的问题中，可用它来进行重采样。

构建WeightedRandomSampler时需提供两个参数：每个样本的权重weights、共选取的样本总数num_samples，以及一个可选参数replacement。权重越大的样本被选中的概率越大，待选取的样本数目一般小于全部的样本数目。replacement用于指定是否可以重复选取某一个样本，默认为True，即允许在一个epoch中重复采样某一个数据。如果设为False，则当某一类的样本被全部选取完，但其样本数目仍未达到num_samples时，sampler将不会再从该类中选择数据，此时可能导致weights参数失效。下面举例说明。

dataset = DogCat('data/dogcat/', transforms=transform)

# 狗的图片被取出的概率是猫的概率的两倍
# 两类图片被取出的概率与weights的绝对大小无关，只和比值有关
weights = [2 if label == 1 else 1 for data, label in dataset]
weights

from torch.utils.data.sampler import  WeightedRandomSampler
sampler = WeightedRandomSampler(weights,\
                                num_samples=9,\
                                replacement=True)
dataloader = DataLoader(dataset,
                        batch_size=3,
                        sampler=sampler)
for datas, labels in dataloader:
    print(labels.tolist())

可见猫狗样本比例约为1:2，另外一共只有8个样本，但是却返回了9个，说明肯定有被重复返回的，这就是replacement参数的作用，下面将replacement设为False试试。

sampler = WeightedRandomSampler(weights, 8, replacement=False)
dataloader = DataLoader(dataset, batch_size=4, sampler=sampler)
for datas, labels in dataloader:
    print(labels.tolist())

在这种情况下，num_samples等于dataset的样本总数，为了不重复选取，sampler会将每个样本都返回，这样就失去weight参数的意义了。

从上面的例子可见sampler在样本采样中的作用：如果指定了sampler，shuffle将不再生效，并且sampler.num_samples会覆盖dataset的实际大小，即一个epoch返回的图片总数取决于sampler.num_samples。

 

2. 计算机视觉工具包：torchvision

计算机视觉是深度学习中最重要的一类应用，为了方便研究者使用，PyTorch团队专门开发了一个视觉工具包torchvion，这个包独立于PyTorch，需通过pip instal torchvision安装。在之前的例子中我们已经见识到了它的部分功能，这里再做一个系统性的介绍。torchvision主要包含三部分：

1）models：提供深度学习中各种经典网络的网络结构以及预训练好的模型，包括AlexNet、VGG系列、ResNet系列、Inception系列等。

2）datasets： 提供常用的数据集加载，设计上都是继承torch.utils.data.Dataset类，主要包括MNIST、CIFAR10/100、ImageNet、COCO等。

3）transforms：提供常用的数据预处理操作，主要包括对Tensor以及PIL Image对象的操作。

from torchvision import models
from torch import nn
# 加载预训练好的模型，如果不存在会进行下载
# 预训练好的模型保存在 ~/.torch/models/下面
#pretrained=True 使用预训练参数  pretrained=False只加载网络结构随机初始化参数
resnet34 = models.squeezenet1_1(pretrained=True, num_classes=1000)

# 修改最后的全连接层为10分类问题（默认是ImageNet上的1000分类）
#resnet34最后一个全联接层的输入是512维的
resnet34.fc=nn.Linear(512, 10)

Transforms中涵盖了大部分对Tensor和PIL Image的常用处理，这些已在上文提到，这里就不再详细介绍。需要注意的是转换分为两步，第一步：构建转换操作，例如transf = transforms.Normalize(mean=x, std=y)，第二步：执行转换操作，例如output = transf(input)。另外还可将多个处理操作用Compose拼接起来，形成一个处理转换流程。

from torchvision import transforms 
to_pil = transforms.ToPILImage() #构建转换操作
to_pil(torch.randn(3, 64, 64)) #执行转换操作 tensor->PIL Image

 

from torchvision import datasets
# 指定数据集路径为data，如果数据集不存在则进行下载
# 通过train=False获取测试集
dataset = datasets.MNIST('data/', download=True, train=False, transform=transform)

len(dataset) #dataset.__len__()

 torchvision还提供了两个常用的函数。一个是make_grid，它能将多张图片拼接成一个网格中；另一个是save_img，它能将Tensor保存成图片。

dataloader = DataLoader(dataset, shuffle=True, batch_size=16)
from torchvision.utils import make_grid, save_image
dataiter = iter(dataloader)
img = make_grid(next(dataiter)[0], 4) # 拼成4*4网格图片，且会转成３通道
to_img(img)

save_image(img, 'a.png')
Image.open('a.png')

 

3. 可视化工具

在训练神经网络时，我们希望能更直观地了解训练情况，包括损失曲线、输入图片、输出图片、卷积核的参数分布等信息。这些信息能帮助我们更好地监督网络的训练过程，并为参数优化提供方向和依据。最简单的办法就是打印输出，但其只能打印数值信息，不够直观，同时无法查看分布、图片、声音等。在本节，我们将介绍两个深度学习中常用的可视化工具：Tensorboard和Visdom。

• Tensorboard

Tensorboard最初是作为TensorFlow的可视化工具迅速流行开来。作为和TensorFlow深度集成的工具，Tensorboard能够展现你的TensorFlow网络计算图，绘制图像生成的定量指标图以及附加数据。但同时Tensorboard也是一个相对独立的工具，只要用户保存的数据遵循相应的格式，tensorboard就能读取这些数据并进行可视化。这里我们将主要介绍如何在PyTorch中使用tensorboardXhttps://github.com/lanpa/tensorboardX进行训练损失的可视化。 TensorboardX是将Tensorboard的功能抽取出来，使得非TensorFlow用户也能使用它进行可视化，几乎支持原生TensorBoard的全部功能。下图为tensorboard界面：

tensorboard的安装主要分为以下两步：

1）安装TensorFlow：如果电脑中已经安装完TensorFlow可以跳过这一步，如果电脑中尚未安装，建议安装CPU-Only的版本，具体安装教程参见TensorFlow官网，或使用pip直接安装，推荐使用清华的软件源。

2）安装tensorboard: pip install tensorboard

3）安装tensorboardX：可通过pip install tensorboardX命令直接安装。

tensorboardX的使用非常简单。首先用如下命令启动tensorboard：

tensorboard --logdir <your/running/dir> --port <your_bind_port>

from tensorboardX import SummaryWriter

# 构建logger对象，logdir用来指定log文件的保存路径
# flush_secs用来指定刷新同步间隔
logger = SummaryWriter(log_dir='experimient_cnn', flush_secs=2)

for ii in range(100):
    logger.add_scalar('data/loss', 10-ii**0.5)
    logger.add_scalar('data/accuracy', ii**0.5/10)

 

在浏览器输入http://appleiMac.local:6006/,即可看到下图的结果：

 

左侧的Horizontal Axis下有三个选项，分别是：

1）Step：根据步长来记录，log_value时如果有步长，则将其作为x轴坐标描点画线。

2）Relative：用前后相对顺序描点画线，可认为logger自己维护了一个step属性，每调用一次log_value就自动加１。

3）Wall：按时间排序描点画线。

左侧的Smoothing条可以左右拖动，用来调节平滑的幅度。点击右上角的刷新按钮可立即刷新结果，默认是每30s自动刷新数据。可见tensorboard_logger的使用十分简单，但它只能统计简单的数值信息，不支持其它功能。

感兴趣的读者可以从github项目主页获取更多信息，本节将把更多的内容留给另一个可视化工具：Visdom。: https://github.com/lanpa/tensorboard-pytorch

 

• Visdom

Visdomhttps://github.com/facebookresearch/visdom是Facebook专门为PyTorch开发的一款可视化工具，其开源于2017年3月。Visdom十分轻量级，但却支持非常丰富的功能，能胜任大多数的科学运算可视化任务。

Visdom可以创造、组织和共享多种数据的可视化，包括数值、图像、文本，甚至是视频，其支持PyTorch、Torch及Numpy。用户可通过编程组织可视化空间，或通过用户接口为生动数据打造仪表板，检查实验结果或调试代码。

Visdom中有两个重要概念：

1）env：环境。不同环境的可视化结果相互隔离，互不影响，在使用时如果不指定env，默认使用main。不同用户、不同程序一般使用不同的env。

2）pane：窗格。窗格可用于可视化图像、数值或打印文本等，其可以拖动、缩放、保存和关闭。一个程序中可使用同一个env中的不同pane，每个pane可视化或记录某一信息。

如下图所示，当前env共有两个pane，一个用于打印log，另一个用于记录损失函数的变化。点击clear按钮可以清空当前env的所有pane，点击save按钮可将当前env保存成json文件，保存路径位于~/.visdom/目录下。也可修改env的名字后点击fork，保存当前env的状态至更名后的env。

Visdom的安装可通过命令pip install visdom。安装完成后，需通过python -m visdom.server命令启动visdom服务，或通过nohup python -m visdom.server &命令将服务放至后台运行。Visdom服务是一个web server服务，默认绑定8097端口，客户端与服务器间通过tornado进行非阻塞交互。

Visdom的使用有两点需要注意的地方：

1) 需手动指定保存env，可在web界面点击save按钮或在程序中调用save方法，否则visdom服务重启后，env等信息会丢失。

2) 客户端与服务器之间的交互采用tornado异步框架，可视化操作不会阻塞当前程序，网络异常也不会导致程序退出。

 Visdom以Plotly为基础，支持丰富的可视化操作，下面举例说明一些最常用的操作。

#%%sh
# 启动visdom服务器
# nohup python -m visdom.server &
import torch 
import visdom

# 新建一个连接客户端
# 指定env = u'test1'，默认端口为8097，host是‘localhost'
vis = visdom.Visdom(env=u'test1',use_incoming_socket=False)

x = torch.arange(1, 30, 0.01)
y = torch.sin(x)
vis.line(X=x, Y=y, win='sinx', opts={'title': 'y=sin(x)'})

下面逐一分析这几行代码：

vis = visdom.Visdom(env=u'test1')，用于构建一个客户端，客户端除指定env之外，还可以指定host、port等参数。vis作为一个客户端对象，可以使用常见的画图函数，包括：

line：类似Matlab中的plot操作，用于记录某些标量的变化，如损失、准确率等

image：可视化图片，可以是输入的图片，也可以是GAN生成的图片，还可以是卷积核的信息

text：用于记录日志等文字信息，支持html格式

histgram：可视化分布，主要是查看数据、参数的分布

scatter：绘制散点图

bar：绘制柱状图

pie：绘制饼状图

更多操作可参考visdom的github主页.

 

这里主要介绍深度学习中常见的line、image和text操作。

Visdom同时支持PyTorch的tensor和Numpy的ndarray两种数据结构，但不支持Python的int、float等类型，因此每次传入时都需先将数据转成ndarray或tensor。上述操作的参数一般不同，但有两个参数是绝大多数操作都具备的：

1) win：用于指定pane的名字，如果不指定，visdom将自动分配一个新的pane。如果两次操作指定的win名字一样，新的操作将覆盖当前pane的内容，因此建议每次操作都重新指定win。

2) opts：选项，接收一个字典，常见的option包括title、xlabel、ylabel、width等，主要用于设置pane的显示格式。

之前提到过，每次操作都会覆盖之前的数值，但往往我们在训练网络的过程中需不断更新数值，如损失值等，这时就需要指定参数update='append'来避免覆盖之前的数值。而除了使用update参数以外，还可以使用vis.updateTrace方法来更新图，但updateTrace不仅能在指定pane上新增一个和已有数据相互独立的Trace，还能像update='append'那样在同一条trace上追加数据。

# append 追加数据
for ii in range(0, 10):
    # y = x
    x = torch.Tensor([ii])
    y = x
    vis.line(X=x, Y=y, win='polynomial', update='append' if ii>0 else None)
    
# updateTrace 新增一条线
x = torch.arange(0, 9, 0.1)
y = (x ** 2) / 9
vis.line(X=x, Y=y, win='polynomial', name='this is a new Trace',update='new')

 

image的画图功能可分为如下两类：

1) image接收一个二维或三维向量，?×?或3×?×?，前者是黑白图像，后者是彩色图像。

2) images接收一个四维向量?×?×?×?，?可以是1或3，分别代表黑白和彩色图像。可实现类似torchvision中make_grid的功能，将多张图片拼接在一起。images也可以接收一个二维或三维的向量，此时它所实现的功能与image一致。

# 可视化一个随机的黑白图片
vis.image(torch.randn(64, 64).numpy())

# 随机可视化一张彩色图片
vis.image(torch.randn(3, 64, 64).numpy(), win='random2')

# 可视化36张随机的彩色图片，每一行6张
vis.images(torch.randn(36, 3, 64, 64).numpy(), nrow=6, win='random3', opts={'title':'random_imgs'})

vis.text用于可视化文本，支持所有的html标签，同时也遵循着html的语法标准。例如，换行需使用<br>标签，\r\n无法实现换行。下面举例说明。

vis.text(u'''<h1>Hello Visdom</h1><br>Visdom是Facebook专门为<b>PyTorch</b>开发的一个可视化工具，
         在内部使用了很久，在2017年3月份开源了它。
         
         Visdom十分轻量级，但是却有十分强大的功能，支持几乎所有的科学运算可视化任务''',
         win='visdom',
         opts={'title': u'visdom简介' }
        )

 

 

4. 使用GPU加速：cuda

这部分内容在前面介绍Tensor、Module时大都提到过，这里将做一个总结，并深入介绍相关应用。

在PyTorch中以下数据结构分为CPU和GPU两个版本：

1)Tensor

2)nn.Module（包括常用的layer、loss function，以及容器Sequential等）

它们都带有一个.cuda方法，调用此方法即可将其转为对应的GPU对象。注意，tensor.cuda会返回一个新对象，这个新对象的数据已转移至GPU，而之前的tensor还在原来的设备上（CPU）。而module.cuda则会将所有的数据都迁移至GPU，并返回自己。所以module = module.cuda()和module.cuda()所起的作用一致。

nn.Module在GPU与CPU之间的转换，本质上还是利用了Tensor在GPU和CPU之间的转换。nn.Module的cuda方法是将nn.Module下的所有parameter（包括子module的parameter）都转移至GPU，而Parameter本质上也是tensor(Tensor的子类)。

下面将举例说明，这部分代码需要你具有两块GPU设备。

为什么将数据转移至GPU的方法叫做.cuda而不是.gpu，就像将数据转移至CPU调用的方法是.cpu？这是因为GPU的编程接口采用CUDA，而目前并不是所有的GPU都支持CUDA，只有部分Nvidia的GPU才支持。PyTorch未来可能会支持AMD的GPU，而AMD GPU的编程接口采用OpenCL，因此PyTorch还预留着.cl方法，用于以后支持AMD等的GPU。

tensor = torch.Tensor(3, 4)
# 返回一个新的tensor，保存在第1块GPU上，但原来的tensor并没有改变
tensor.cuda(0) #tensor.cuda() 参数默认为0
tensor.is_cuda # False

# 不指定所使用的GPU设备，将默认使用第1块GPU
tensor = tensor.cuda()
tensor.is_cuda # True

module = nn.Linear(3, 4)
module.cuda(device = 1) #等价于module = module.cuda(device=1) 可以指定哪块gpu 默认为0
module.weight.is_cuda # True

class VeryBigModule(nn.Module):
    def __init__(self):
        super(VeryBigModule, self).__init__()
        self.GiantParameter1 = torch.nn.Parameter(torch.randn(100000, 20000)).cuda(0)
        self.GiantParameter2 = torch.nn.Parameter(torch.randn(20000, 100000)).cuda(1)
    
    def forward(self, x):
        x = self.GiantParameter1.mm(x.cuda(0))
        x = self.GiantParameter2.mm(x.cuda(1))
        return x

 上面最后一部分中，两个Parameter所占用的内存空间都非常大，大概是8个G，如果将这两个都同时放在一块GPU上几乎会将显存占满，无法再进行任何其它运算。此时可通过这种方式将不同的计算分布到不同的GPU中。

关于使用GPU的一些建议：

1) GPU运算很快，但对于很小的运算量来说，并不能体现出它的优势，因此对于一些简单的操作可直接利用CPU完成

2) 数据在CPU和GPU之间，以及GPU与GPU之间的传递会比较耗时，应当尽量避免

3) 在进行低精度的计算时，可以考虑HalfTensor，它相比于FloatTensor能节省一半的显存，但需千万注意数值溢出的情况。

另外这里需要专门提一下，大部分的损失函数也都属于nn.Moudle，但在使用GPU时，很多时候我们都忘记使用它的.cuda方法，这在大多数情况下不会报错，因为损失函数本身没有可学习的参数（learnable parameters）。但在某些情况下会出现问题，为了保险起见同时也为了代码更规范，应记得调用criterion.cuda。下面举例说明。

# 交叉熵损失函数，带权重
criterion = torch.nn.CrossEntropyLoss(weight=torch.Tensor([1, 3]))
input = torch.randn(4, 2).cuda()
target = torch.Tensor([1, 0, 0, 1]).long().cuda()

# 下面这行会报错，因weight未被转移至GPU
# loss = criterion(input, target)

# 这行则不会报错
criterion.cuda()
loss = criterion(input, target)

criterion._buffers

而除了调用对象的.cuda方法之外，还可以使用torch.cuda.device，来指定默认使用哪一块GPU，或使用torch.set_default_tensor_type使程序默认使用GPU，不需要手动调用cuda。

# 如果未指定使用哪块GPU，默认使用GPU 0
x = torch.cuda.FloatTensor(2, 3)
# x.get_device() == 0
y = torch.FloatTensor(2, 3).cuda()
# y.get_device() == 0

# 指定默认使用GPU 1
with torch.cuda.device(1):    
    # 在GPU 1上构建tensor
    a = torch.cuda.FloatTensor(2, 3)

    # 将tensor转移至GPU 1
    b = torch.FloatTensor(2, 3).cuda()
    print(a.get_device() == b.get_device() == 1 )

    c = a + b
    print(c.get_device() == 1)

    z = x + y
    print(z.get_device() == 0)

    # 手动指定使用GPU 0
    d = torch.randn(2, 3).cuda(0)
    print(d.get_device() == 2)

 

torch.set_default_tensor_type('torch.cuda.FloatTensor') # 指定默认tensor的类型为GPU上的FloatTensor
a = torch.ones(2, 3)
a.is_cuda

如果服务器具有多个GPU，tensor.cuda()方法会将tensor保存到第一块GPU上，等价于tensor.cuda(0)。此时如果想使用第二块GPU，需手动指定tensor.cuda(1)，而这需要修改大量代码，很是繁琐。这里有两种替代方法：

1) 一种是先调用torch.cuda.set_device(1)指定使用第二块GPU，后续的.cuda()都无需更改，切换GPU只需修改这一行代码。

2) 更推荐的方法是设置环境变量CUDA_VISIBLE_DEVICES，例如当export CUDA_VISIBLE_DEVICE=1（下标是从0开始，1代表第二块GPU），只使用第二块物理GPU，但在程序中这块GPU会被看成是第一块逻辑GPU，因此此时调用tensor.cuda()会将Tensor转移至第二块物理GPU。CUDA_VISIBLE_DEVICES还可以指定多个GPU，如export CUDA_VISIBLE_DEVICES=0,2,3，那么第一、三、四块物理GPU会被映射成第一、二、三块逻辑GPU，tensor.cuda(1)会将Tensor转移到第三块物理GPU上。

设置CUDA_VISIBLE_DEVICES有两种方法，一种是在命令行中CUDA_VISIBLE_DEVICES=0,1 python main.py，一种是在程序中import os;os.environ["CUDA_VISIBLE_DEVICES"] = "2"。如果使用IPython或者Jupyter notebook，还可以使用%env CUDA_VISIBLE_DEVICES=1,2来设置环境变量。

从 0.4 版本开始，pytorch新增了tensor.to(device)方法，能够实现设备透明，便于实现CPU/GPU兼容。这部份内容已经在上篇博客中讲解过了。

从PyTorch 0.2版本中，PyTorch新增分布式GPU支持。分布式是指有多个GPU在多台服务器上，而并行一般指的是一台服务器上的多个GPU。分布式涉及到了服务器之间的通信，因此比较复杂，PyTorch封装了相应的接口，可以用几句简单的代码实现分布式训练。分布式对普通用户来说比较遥远，因为搭建一个分布式集群的代价十分大，使用也比较复杂。相比之下一机多卡更加现实。对于分布式训练，这里不做太多的介绍，感兴趣的读者可参考文档http://pytorch.org/docs/distributed.html。

• 单机多卡并行

要实现模型单机多卡十分容易，直接使用 new_module = nn.DataParallel(module, device_ids), 默认会把模型分布到所有的卡上。多卡并行的机制如下：

1）将模型（module）复制到每一张卡上

2）将形状为（N,C,H,W）的输入均等分为 n份（假设有n张卡），每一份形状是（N/n, C,H,W）,然后在每张卡前向传播，反向传播，梯度求平均。要求batch-size 大于等于卡的个数(N>=n)

在绝大多数情况下，new_module的用法和module一致，除了极其特殊的情况下（RNN中的PackedSequence）。另外想要获取原始的单卡模型，需要通过new_module.module访问。

• 多机分布式 

5. 持久化

在PyTorch中，以下对象可以持久化到硬盘，并能通过相应的方法加载到内存中：

1）Tensor、Variable

2）nn.Module

3）Optimizer

本质上上述这些信息最终都是保存成Tensor。Tensor的保存和加载十分的简单，使用torch.save和torch.load即可完成相应的功能。在save/load时可指定使用的pickle模块，在load时还可将GPU tensor映射到CPU或其它GPU上。

我们可以通过torch.save(obj, file_name)等方法保存任意可序列化的对象，然后通过obj = torch.load(file_name)方法加载保存的数据。对于Module和Optimizer对象，这里建议保存对应的state_dict，而不是直接保存整个Module/Optimizer对象。Optimizer对象保存的主要是参数，以及动量信息，通过加载之前的动量信息，能够有效地减少模型震荡，下面举例说明。

a = torch.Tensor(3, 4)
if torch.cuda.is_available():
        a = a.cuda(1) # 把a转为GPU1上的tensor,
        torch.save(a,'a.pth')
        
        # 加载为b, 存储于GPU1上(因为保存时tensor就在GPU1上)
        b = torch.load('a.pth')
        
        # 加载为c, 存储于CPU
        c = torch.load('a.pth', map_location=lambda storage, loc: storage)
        
        # 加载为d, 存储于GPU0上
        d = torch.load('a.pth', map_location={'cuda:1':'cuda:0'})

torch.set_default_tensor_type('torch.FloatTensor')
from torchvision.models import SqueezeNet
model = SqueezeNet()
# module的state_dict是一个字典
model.state_dict().keys()

# Module对象的保存与加载
torch.save(model.state_dict(), 'squeezenet.pth')
model.load_state_dict(torch.load('squeezenet.pth'))

optimizer = torch.optim.Adam(model.parameters(), lr=0.1)
torch.save(optimizer.state_dict(), 'optimizer.pth')
optimizer.load_state_dict(torch.load('optimizer.pth'))

all_data = dict(
    optimizer = optimizer.state_dict(),
    model = model.state_dict(),
    info = u'模型和优化器的所有参数'
)
torch.save(all_data, 'all.pth')
all_data = torch.load('all.pth')
all_data.keys()