【notes】pytorch学习笔记4-pytorch常用工具

发表于 分类于 ，
本文字数： 17k 阅读时长 ≈ 15 分钟

pytorch中常见的工具

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

数据处理

在解决深度学习问题的过程中，往往需要花费大量的精力去处理数据，包括图像、文本、语音或其他二进制数据等。数据的处理对训练神经网络来说十分重要，良好的数据处理不仅会加速模型训练，也会提高模型效果。

数据加载

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

  • __getitem__：返回一条数据或一个样本。obj[index]等价于obj.__getitem__(index)。
  • __len__：返回样本的数量。len(obj)等价于obj.__len__()。

这里我们用Kaggle经典挑战赛“Dogs vs.Cat”

1
2
%env LS_COLORS = None
!tree --charset ascii data/dogcat/
env: LS_COLORS=None
data/dogcat/
|-- cat.0.jpg
|-- cat.1.jpg
|-- cat.2.jpg
|-- cat.3.jpg
|-- cat.4.jpg
|-- cat.5.jpg
|-- cat.6.jpg
|-- dog.0.jpg
|-- dog.1.jpg
|-- dog.10.jpg
|-- dog.1000.jpg
`-- dog.10000.jpg

0 directories, 12 files
1
2
3
import torch as t
from torch.utils import data
import matplotlib.pyplot as plt
1
2
3
import os
from PIL import Image
import numpy as np
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
class DogCat(data.Dataset):
    def __init__(self, root):
        imgs = os.listdir(root)
        # 所有图片的绝对路径
        self.imgs = [os.path.join(root, img) for img in imgs]
    
    def __getitem__(self, 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 = t.from_numpy(array)
        return data, label
    def __len__(self):
        return len(self.imgs)
1
2
3
4
dataset = DogCat('./data/dogcat/')
img, label = dataset[0]
for img, label in dataset:
    print(img.size(), img.float().mean(), label)
torch.Size([375, 499, 3]) tensor(116.7904) 1
torch.Size([499, 327, 3]) tensor(133.5602) 1
torch.Size([144, 175, 3]) tensor(166.6151) 0
torch.Size([292, 269, 3]) tensor(157.4856) 1
torch.Size([375, 499, 3]) tensor(96.8243) 0
torch.Size([375, 499, 3]) tensor(120.7302) 1
torch.Size([280, 300, 3]) tensor(71.6653) 0
torch.Size([396, 312, 3]) tensor(131.8400) 0
torch.Size([303, 400, 3]) tensor(129.1319) 0
torch.Size([374, 500, 3]) tensor(119.7826) 0
torch.Size([412, 263, 3]) tensor(152.9542) 1
torch.Size([414, 500, 3]) tensor(156.6921) 0

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

  • 返回样本的形状之一，每张图片的大小不一样，这对于需要去batch训练的神经网络来说很不友好。
  • 返回样本的数值较大，为归一化至[-1, 1]
    针对上述问题，pytorch提供了torchvision。它是一个视觉工具包，提供了很多视觉图像处理的工具，其中transforms模块提供了对PIL Image对象和Tensor对象的常用操作。

针对上述问题，pytorch提供torchvision。它是一个视觉工具包，提供了很多视觉图像处理的工具，其中transforms模块提供了对PIL Image对象和Tensor对象的常用操作。
对PIL Image的常见操作如下。

  • Resize：调整图片尺寸
  • CenterCrop、RandomCrop、RandomSizedCrop：裁剪图片。
  • Pad：填充。
  • ToTensor：将PIL Image对象转成Tensor，会自动将[0, 225]归一化至[0, 1]。

对Tensor的常见操作如下。

  • Normalize：标准化，即减均值，除以标准差。
  • ToPILImage：将Tensor转为PIL Image对象。

如果要对图片进行多个操作，可通过Compose将这些操作拼接起来，类似于nn.Sequential。注意，这些操作定义后是以对象的形式存在，真正使用时需要调用它的__call__方法，类似于nn.Module。下面利用这些操作来优化上面的dataset。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
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 x224的图片
    T.ToTensor(),   # 将图片转换成Tensor并归一化至[0, 1]
    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)
1
2
3
4
5
from torchvision.transforms import ToPILImage
dataset = DogCat('./data/dogcat/', transforms = transform)
img, label = dataset[0]
for img, label in dataset:
    print(img.size(), label)
torch.Size([3, 224, 224]) 0
torch.Size([3, 224, 224]) 0
torch.Size([3, 224, 224]) 1
torch.Size([3, 224, 224]) 0
torch.Size([3, 224, 224]) 1
torch.Size([3, 224, 224]) 0
torch.Size([3, 224, 224]) 1
torch.Size([3, 224, 224]) 1
torch.Size([3, 224, 224]) 1
torch.Size([3, 224, 224]) 1
torch.Size([3, 224, 224]) 0
torch.Size([3, 224, 224]) 1

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

torchvision已经预先实现了常用的Dataset，包括前面使用过的CIFAR-10，以及ImageNet、COCO、MNIST、LSUN等数据集，可通过调用torchvision.datasets下相应对象来调用相关数据集。
还有一个经常使用到的Dataset——ImageFolder，它的实现和上述DogCat很相似。ImageFolder假设所有的文件按文件夹保存，每个文件夹下储存同一个类别的图片，文件夹名为类名，其构造函数如下：

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

它主要有以下四个参数。

  • root：在root指定的路径下寻找图片。
  • transform：对PIL Image进行转换操作，transform的输入是使用loader读取图片的返回对象。
  • target_transform：对label的转换。
  • loader：指定加载图片的函数，默认操作是读取为PIL Image对象。

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

1
!tree --charset ASCII data/dogcat2/
data/dogcat2/
|-- cat
|   |-- cat.0.jpg
|   |-- cat.1.jpg
|   |-- cat.2.jpg
|   |-- cat.3.jpg
|   `-- cat.4.jpg
`-- dog
    |-- dog.0.jpg
    |-- dog.1.jpg
    |-- dog.10.jpg
    |-- dog.1000.jpg
    `-- dog.10000.jpg

2 directories, 10 files
1
2
from torchvision.datasets import ImageFolder
dataset = ImageFolder('data/dogcat2/')
1
dataset.class_to_idx
{'cat': 0, 'dog': 1}
1
dataset.imgs
[('data/dogcat2/cat/cat.0.jpg', 0),
 ('data/dogcat2/cat/cat.1.jpg', 0),
 ('data/dogcat2/cat/cat.2.jpg', 0),
 ('data/dogcat2/cat/cat.3.jpg', 0),
 ('data/dogcat2/cat/cat.4.jpg', 0),
 ('data/dogcat2/dog/dog.0.jpg', 1),
 ('data/dogcat2/dog/dog.1.jpg', 1),
 ('data/dogcat2/dog/dog.10.jpg', 1),
 ('data/dogcat2/dog/dog.1000.jpg', 1),
 ('data/dogcat2/dog/dog.10000.jpg', 1)]
1
2
3
# 没有任何transform，所以返回的还是PIL Image对象
dataset[0][1] # 第一维是第几张图，第二维为1是label，为0是Image对象
dataset[0][0]

png

1
2
3
4
5
6
7
8
# 加上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,
])
1
dataset = ImageFolder('data/dogcat2/', transform = transform)
1
dataset.class_to_idx
1
2
3
# 深度学习中图片数据一般保存为CHW，即通道数 x 高 x 宽
for img, index in dataset:
    print(img.size(), index)
torch.Size([3, 224, 224]) 0
torch.Size([3, 224, 224]) 0
torch.Size([3, 224, 224]) 0
torch.Size([3, 224, 224]) 0
torch.Size([3, 224, 224]) 0
torch.Size([3, 224, 224]) 1
torch.Size([3, 224, 224]) 1
torch.Size([3, 224, 224]) 1
torch.Size([3, 224, 224]) 1
torch.Size([3, 224, 224]) 1
1
2
3
to_img = T.ToPILImage()
# 0.2和0.4是标准差和均值的近似
to_img(dataset[0][0]*0.2 + 0.4)

png

Dataset只负责数据的抽象，一次调用__getattr__只返回一个样本。而训练神经网络是对一个batch的数据进行操作，同时还需要对数据进行shuffle和并行加速等。对此，pytorch提供了DataLoader帮助我们实现这些功能。
DataLoader的函数定义如下。

1
DataLoader(dataset, batch_size = 1, shuffle = False, sampler = None, num_workers = 0, collate_fn = default_collate, pin_memory = False, drop_last = False)
  • dataset：加载的数据集（Dataset对象）
  • batch_size：batch size（批大小）
  • shuffle：是否将数据打乱
  • sampler：样本抽样，后续会详细介绍。
  • num_workers：使用多进程加载的进程数，0代表不使用多进程。
  • collate_fn：如何将多个样本数据拼接成一个batch，一般使用默认的拼接方式即可。
  • pin_memory：是否将数据保存在pin memory区，pin memory中的数据转到GPU会快一些。
  • drop_last：dataset中的数据个数可能不是batch_size的整数倍，drop_last为True会将多出来不足一个batch的数据丢失。
1
from torch.utils.data import DataLoader
1
dataloader = DataLoader(dataset, batch_size = 3, shuffle = True, num_workers = 0, drop_last = False)
1
2
3
dataiter = iter(dataloader)
imgs, labels = next(dataiter)
imgs.size()
torch.Size([3, 3, 224, 224])

dataloader是一个可迭代对象，我们可以向使用迭代器一样使用它，利用：

1
2
for batch_datas, batch_labels in dataloader:
    train()

或

1
2
dataiter = iter(dataloader)
batch_datas, batch_labels = next(dataiter)

在数据处理中，有时会出现某个样本无法读取等问题，例如某张图片损坏。这时__getitem__函数中将出现异常，此时最好的解决方案即是将出错的样本剔除。如果遇到这种情况实在无法处理，则可以返回None对象，然后在Dataloader中实现自定义的collate_fn，将空对象过滤掉。但要注意，在这种情况下dataloader返回的一个batch的样本数目会少于batch_size。
对丢弃样本异常图片而言，这种做法会更好一些，因为它能保证每个batch样本的数目仍是batch_size。但在大多数情况下，最好的方式还是对数据进行彻底清除。

DataLoader里并没有太多的魔法方法，它封装了python的标准库multiprocessing使其能够实现多进程加速。在Dataset和DataLoader的使用方面有以下建议。

  1. 高负载的操作放在__getitem__中，如加载图片
  2. dataset中应尽量只包含只读对象，避免修改任何可变对象。

第一点是因为多进程会并行地调用__getitem__函数，将负载高的放在__getitem__函数中能够实现并行加速。第二点是因为dataloader使用多进程加载，如果在Dataset中使用了可变对象，可能会有意想不到的冲突。
在多线程/多进程中，修改一个可变对象需要加锁，但是dataloader的设计使得其很难加锁（在实际使用中也应尽量避免锁的存在）。
如果一定要修改可变对象，建议使用python标准库queue
使用python multiprocessing库的另一个问题是，在使用多进程时，如果主程序异常终止（比如用“Ctrl+C”快捷键强行退出），相应的数据加载进程可能无法正常退出。这时需要手动强行终止进程。

1
ps x | grep <cmdline> | awk '{print $1}' | xargs kill
  • ps x:获取当前用户的所有进程。
  • grep : 找到已经停止的pytorch程序的进程，例如你是通过python train.py启动的，那就需要些grep ‘python train.py’。
  • awk ‘{print $1}’：获取进程的pid
  • xargs kill：终止进程，根据需要可能要写成xargs kill -9强制终止进程。

pytorch还单独提供一个sampler模块，用来对数据进行采样。常用的有随机采样器RandomSampler，当dataloader的shuffle参数为True时，就是调用的这个。
这里介绍一个很有用的采样方法：WeightedRandomSampler，它会根据每个样本的权重选取数据，在样本比例不均衡的问题中，可用它进行重采样。

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

1
!tree --charset ASCII data/dogcat/
data/dogcat/
|-- cat.0.jpg
|-- cat.1.jpg
|-- cat.2.jpg
|-- cat.5.jpg
|-- cat.6.jpg
|-- dog.0.jpg
|-- dog.1.jpg
|-- dog.10.jpg
|-- dog.1000.jpg
`-- dog.10000.jpg

0 directories, 10 files
1
2
3
4
5
6
dataset = DogCat('data/dogcat',transforms = transform)

# 狗的图片被取出的概率是猫的概率的两倍
# 两类图片被取出的概率与weights的绝对大小无关，只和比值有关
weights = [2 if label == 1 else 1 for data, label in dataset]
weights
[1, 1, 2, 1, 1, 2, 2, 2, 2, 1]
1
2
3
4
5
from torch.utils.data.sampler import WeightedRandomSampler
sampler = WeightedRandomSampler(weights, num_samples=12, replacement=True)
dataloader = DataLoader(dataset, batch_size = 3, sampler=sampler)
for datas, labels in dataloader:
    print(labels.tolist())
[0, 0, 1]
[0, 1, 1]
[1, 1, 1]
[0, 0, 0]

一共只有10个样本，却返回了12个，说明样本被重复返回，这就是replacement参数的作用

1
2
3
4
5
# replacement改为False
sampler = WeightedRandomSampler(weights, num_samples=10, replacement=False)
dataloader = DataLoader(dataset, batch_size = 5, sampler=sampler)
for datas, labels in dataloader:
    print(labels.tolist())
[0, 1, 0, 1, 1]
[1, 1, 0, 0, 0]

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

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

计算机视觉工具包：torchvision

计算机视觉是深度学习中最重要的一类应用，为了方便研究者使用，pytorch团队专门开发一个视觉工具包torchvision，这个包独立于pytorch。

torchvision主要包含以下三部分。

  • model：提供深度学习中各种经典网络的网络结构及预训练好的模型，包括AlexNet、VGG系列、ResNet系列、Inception系列等。
  • datasets：提供常用的数据集加载，设计上都是继承torch.utils.data.Dataset，主要包括MNIST、CIFAR10/100、ImageNet、coco等。
  • transforms：提供常用的数据预处理操作，主要包括对Tensor及PIL Image对象的操作。
1
2
3
4
5
6
7
8
9
from torchvision import models
from torch import nn

# 加载预训练好的模型，如果不存在会下载
# 预训练好的模型保存在 ~/.torch/models/下面
resnet34 = models.resnet34(pretrained = True,num_classes = 1000)

# 修改最后的全连接层为10分类问题（默认是ImageNet上的1000分类）
resnet34.fc = nn.Linear(512, 10)
Downloading: "https://download.pytorch.org/models/resnet34-333f7ec4.pth" to /root/.cache/torch/checkpoints/resnet34-333f7ec4.pth
100%|██████████| 83.3M/83.3M [00:35<00:00, 2.49MB/s]
1
2
3
4
5
6
transform = T.Compose([
    T.Resize(32), # 缩放图片（Image），保持长宽比不变，最短边为224像素
    T.CenterCrop(32), # 从中间取出224 x224的图片
    T.ToTensor(),   # 将图片转换成Tensor并归一化至[0, 1]
    T.Normalize(mean = [.5], std = [.5,]) # 标准化至[-1, 1]
])
1
2
3
4
from torchvision import datasets
# 指定数据集路径为data，如果数据集不存在则进行下载
# 通过train = False获取测试集
dataset = datasets.MNIST('data/',download=True, train = False, transform = transform)

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

1
len(dataset)
10000
1
2
3
4
5
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网络图片，且会转成3通道
to_img(img)

png

1
2
3
img2 = make_grid(next(dataiter)[0], 4) # 拼成4*4网络图片，且会转成3通道
save_image(img2, 'a.png')
Image.open('a.png')

png

可视化工具

visdom

visdom是facebook专门为pytorch开发的一款可视化工具，开源于2017年3月。

visdom可以创造、组织和共享多种数据的可视化，包括数值、图像、文本，甚至是视频，支持pytorch、torch及numpy。用户可通过编程组织可视化空间或通过用户接口为数据打造仪表盘，检查试验结果和调试代码。
visdom中有一下两个重要概念。

  • env：环境。不同环境的可视化结果相互隔离，互不影响，在使用时如果不指定env，默认使用main。不同用户。不同程序一般使用不同的env。
  • pane：窗格。窗格可用于可视化图像、数值或打印文本等，其可以拖动、缩放、保存和关闭。一个程序可使用同一个env中的不同pane，每个pane可视化或记录某一信息。
    “clear”按钮可以清空当前env的所有pane，“save”按钮可将当前env保存成json文件，保存路径位于~/.visdom/目录下。修改env的名字后单击fork，可将当前env另存为新文件。

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

使用visdom时有两点需要注意的地方。

  • 需手动指定保存env，可在web界面单击save按钮或在程序中调用save方法，否则visdom服务重启后，env等信息会丢失
  • 客户端与服务器之间的交互采用tornado异步框架，可视化操作不会阻塞当前程序，网络异常也不会导致程序退出。
    visdom以Plotly为基础。
1
2
3
4
5
import visdom
vis = visdom.Visdom(env = u'test1')
x = t.arange(1, 30, 0.01)
y = t.sin(x)
vis.line(X = x, Y = y, win = 'sinx',opts = {'title':'y = sin(x)'})
WARNING:root:Setting up a new session...





'sinx'

下面我们逐一分析这几行代码。

  • vis = visdom.Visdom(env = u’test1’)，用于构建一个客户端，客户端除了制定env外，还可以制定host、post等参数。
  • vis作为一个客户端对象，可以使用如下常见的画图函数。
    • line：类似MATLAB中的plot操作，用于记录某些标量的变化，例如损失、标准率等。
    • image：可视化图片，可以是输入的图片，也可以是GAN生成的图片，还可以是卷积核的信息。
    • text：用于记录日志等文字信息，支持HTML格式
    • histgram：可视化分布，主要是查看数据、参数的分布。
    • scatter：绘制散点图。
    • bar：绘制柱状图
    • pie：绘制饼状图
    • 更多参考github主页

visdom同时支持pytorch的tensor和numpy的ndarray两种数据结构，但不支持python的int、float等类型。上述操作的参数一般不同，但有两个参数是绝大多数操作都具备。

  • win：用于指定pane的名字，如果不指定，visdom将自动分配一个新的pane。如果两次操作指定的win名字一样，新的操作将覆盖当前pane的内容，因此建议每次操作都指定win
  • opts：用来可视化配置，接受一个字典，常见的option包括title、xlabel、ylabel、width等，主要用于设置pane的显示格式。

之前提到过，每次操作会覆盖之前的数据，但我们在训练网络的过程中往往需要不断更新数值，这时就需要指定参数update=’append’来避免覆盖之前的数值。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
import time
# append追加数据
for ii in range(0, 10):
    x = t.Tensor([ii])
    y = 2 * x
    time.sleep(0.5)
    vis.line(X = x, Y= y, win = 'polynomial',name = 'Trace1', update = 'append' if ii > 0 else None)
    
# 新增一条线
for ii in np.linspace(0, 5, 50):
    x = t.Tensor([ii])
    y = x ** 2
    time.sleep(0.1)
    vis.line(X = x, Y= y, win = 'polynomial',name = 'Trace2', update = 'append')

images的画图功能可分为如下两类。

  • image接受一个二维或三维向量，HW 或 3H*W ，前者是黑白图像，后者是彩色图像。
  • images接受一个四维向量NCH*W，C可以是1或3，分别代表黑白和彩色图像。可实现类似torchvision中make_grid的功能，将多张图片拼接在一起。images也可以接受一个二维或三维的向量，此时它所实现的功能与image一致。
1
2
3
4
5
6
7
8
# 可视化一张随机的黑白照片
vis.image(t.randn(64, 64).numpy())

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

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

x = dataset[0][0].unsqueeze(0)
for i in range(1,36):
    x = t.cat((x, dataset[i][0].unsqueeze(0)), dim = 0)
vis.images(x, nrow = 6, win = 'MNIST',opts = {'title':'MNIST数据集'})
'MNIST'

vis.text用于可视化文本，它支持所有的html标签，同时也遵循着html的语法标签。

1
vis.text(u'''<h1>Hello visdom</h1><br>visdom是Facebook专门为<b>pytorch<b/>开发一个可视化工具，''',win = 'visdom',opts = {'title':u'visdom简介'})
'visdom'
1
vis.text('ss', win = 'visdom', append = True, opts = {'title':'平方'})
'visdom'
1
2
3
4
5
# 文本更新
vis.text('<b>平方:</b>', win = 'pingfang',opts = {'title':'pingfang'})
for i in range(1, 20):
    time.sleep(0.5)
    vis.text('<b>[info]{}^2={}</b>'.format(i, i **2), append = True,win = 'pingfang')

GPU

略

持久化

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

  • Tensor
  • Variable
  • nn.Module
  • Optimizer

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

我们可以通过t.save(obj, file_name)等方法保存任意可序列化的对象，然后通过obj = t.load(file_name)方法加载保存的数据。

对Module和Optimizer对象，这里建议保存对应的state_dict，而不是直接保存整个Module/Optimizer对象。Optimizer对象保存的是参数即动量信息，通过加载之前的动量信息，能够有效减少模型震荡。

1
2
3
4
5
6
7
8
9
10
11
12
13
a = t.Tensor(3, 4)
if t.cuda.is_available():
    a = a.cuda(1)# 把a转为GPU1上的tensor
    t.save(a, 'a.pth')
    
    #加载为b，存储于GPU1上（因为保存时就在GPU1）
    b = t.load('a.pth')
    
    # 加载为c，存储于CPU
    c = t.load('a.pth', map_location = lambda storage, loc:storage)
    
    # 加载为d，存储于GPU0
    d = t.load('a.pth', map_location = {'cuda:1':'cuda:0'})
1
2
3
from torchvision.models import resnet34

model = resnet34()
1
2
# model的state_dict是一个字典
list(model.state_dict().keys())[:8]
['conv1.weight',
 'bn1.weight',
 'bn1.bias',
 'bn1.running_mean',
 'bn1.running_var',
 'bn1.num_batches_tracked',
 'layer1.0.conv1.weight',
 'layer1.0.bn1.weight']
1
2
3
# module对象的保存与加载
t.save(model.state_dict(), 'resnet32.pth')
model.load_state_dict(t.load('resnet32.pth'))
<All keys matched successfully>
1
optimizer = t.optim.Adam(model.parameters(), lr = 0.1)
1
2
t.save(optimizer.state_dict(),'optimizer.pth')
optimizer.load_state_dict(t.load('optimizer.pth'))
1
2
3
4
5
6
all_data = dict(
    optimizer = optimizer.state_dict(),
    model = model.state_dict(),
    info = u'模型和优化器的所有参数'
)
t.save(all_data, 'all_data.pth')
1
2
all_data = t.load('all_data.pth')
all_data.keys()
dict_keys(['optimizer', 'model', 'info'])

———————————————感谢阅读———————————————

欢迎收藏访问我的博客 知乎 掘金 简书 知乎

相关文章
贰三 wechat
欢迎扫描二维码订阅我的公众号！