PyTorch 笔记（12）— Tensor 持久化、向量化、torch.set_num_threads、torch.set_printoptions

1. 持久化

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

• Tensor
• Variable
• nn.Module
• Optimizer

本质上上述这些信息最终都是保存成 Tensor 。

Tensor 的保存和加载也比较简单，使用 t.save 个 t.load 即可完成相应的功能，在 save、load时可指定使用的 pickle 模块，在 load 时还可将 GPU Tensor 映射到 CPU 或其它 GPU 上。

In [3]: import torch as t

In [4]: a = t.arange(0,10)

In [5]: if t.cuda.is_available():
   ...:     a = a.cuda(1)
   ...:     t.save("a.pth")
   ...:     b = t.load("a.pth")
   ...:     c = t.load("a.pth", map_location=lambda storage, loc:storage)
   ...:     d = t.load("a.pth", map_location={'cuda:1':'cuda:0'}) 

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

a = t.Tensor(3, 4)
if t.cuda.is_available():
        a = a.cuda(0) # 把a转为GPU0上的tensor,
        t.save(a,'a.pth')
        
        # 加载为b, 存储于GPU0上(因为保存时tensor就在GPU0上)
        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'})

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

输出：

odict_keys(['features.0.weight', 'features.0.bias', 'features.3.squeeze.weight', 'features.3.squeeze.bias', 'features.3.expand1x1.weight', 'features.3.expand1x1.bias', 'features.3.expand3x3.weight', 'features.3.expand3x3.bias', 'features.4.squeeze.weight', 'features.4.squeeze.bias', 'features.4.expand1x1.weight', 'features.4.expand1x1.bias', 'features.4.expand3x3.weight', 'features.4.expand3x3.bias', 'features.5.squeeze.weight', 'features.5.squeeze.bias', 'features.5.expand1x1.weight', 'features.5.expand1x1.bias', 'features.5.expand3x3.weight', 'features.5.expand3x3.bias', 'features.7.squeeze.weight', 'features.7.squeeze.bias', 'features.7.expand1x1.weight', 'features.7.expand1x1.bias', 'features.7.expand3x3.weight', 'features.7.expand3x3.bias', 'features.8.squeeze.weight', 'features.8.squeeze.bias', 'features.8.expand1x1.weight', 'features.8.expand1x1.bias', 'features.8.expand3x3.weight', 'features.8.expand3x3.bias', 'features.9.squeeze.weight', 'features.9.squeeze.bias', 'features.9.expand1x1.weight', 'features.9.expand1x1.bias', 'features.9.expand3x3.weight', 'features.9.expand3x3.bias', 'features.10.squeeze.weight', 'features.10.squeeze.bias', 'features.10.expand1x1.weight', 'features.10.expand1x1.bias', 'features.10.expand3x3.weight', 'features.10.expand3x3.bias', 'features.12.squeeze.weight', 'features.12.squeeze.bias', 'features.12.expand1x1.weight', 'features.12.expand1x1.bias', 'features.12.expand3x3.weight', 'features.12.expand3x3.bias', 'classifier.1.weight', 'classifier.1.bias'])

# Module对象的保存与加载
t.save(model.state_dict(), 'squeezenet.pth')
model.load_state_dict(t.load('squeezenet.pth'))	# <All keys matched successfully>

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

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

all_data = t.load('all.pth')
all_data.keys() 

输出：

dict_keys(['optimizer', 'model', 'info'])

2. 向量化

向量化计算是一种特殊的并行计算方式，一般程序在同一时间只执行一个操作指令，而向量化计算可在同一时间执行多个操作，通常是对不同的数据执行同样的一个或一批指令，或者说把指令应用于一个数组或向量上。

向量化可极大地提高科学运算效率，在科学计算中应当极力避免使用 Python 原生的 for 循环，尽量使用向量化的数值运算。

In [1]: import torch as t

In [2]: def for_loop_add(x, y):
   ...:     result = []
   ...:     for i, j in zip(x, y):
   ...:         result.append(i+j)
   ...:     return t.Tensor(result)
   ...:     

In [3]: x = t.zeros(100)

In [4]: y = t.ones(100)

In [5]: %timeit -n 10 for_loop_add(x, y)
The slowest run took 44.68 times longer than the fastest. This could mean that an intermediate result is being cached.
5.04 ms ± 10.1 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

In [6]: %timeit -n 10 x + y
The slowest run took 59.81 times longer than the fastest. This could mean that an intermediate result is being cached.
25.8 µs ± 52.9 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

可以看出两者运算速度会差很大，因此实际使用中尽量使用内建函数，这些函数底层是由 C/C++ 实现，能通过执行底层优化实现高效运算。

3. Tensor 特点

• 大多数的 t.function 都有一个参数 out ，这时产生的结果会保存在 out 指定的 tensor 之中；
• t.set_num_threads 可以设置 PyTorch 进行 CPU 多线程并行计算时所占用的线程数，用来限制 PyTorch 所占用的 CPU 数目；
• t.set_printoptions可以用来设置打印 tensor 时的数值精度和格式；