torch.device
在每次的使用pytorch的开头我们都要配置好我们训练使用的设备，使用cpu还是gpu
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
to(device)
主要要将两部分加入device：
-
模型model
-
创建的所有的tensor（包括所有输入的数据和标签，一些初始化的状态，如rnn的h0）
使用model.to(device)或tensor.to(device)将model和中间创建的Tensor加入device即可
# 将船创建好的Tensor加入device a = torch.randn(3,1) a = a.to(device) # 或直接创建一个加入device的Tensor a = torch.randn(3,1).to(device) # 将模型加入device model = seq2seq() model = model.to(device) # 很常用：在loader中读取的数据加入device for batch_x, batch_y in loader: batch_x, batch_y = batch_x.to(device), batch_y.to(device)
torch.device使用方法
torch.device代表将torch.Tensor分配到的设备的对象，有cpu和cuda两种，这里的cuda就是gpu，至于为什么不直接用gpu与cpu对应，是因为gpu的编程接口采用的是cuda
print(torch.cuda.is_available())
#cuda是否可用；
print(torch.cuda.device_count())
#返回gpu数量；
print(torch.cuda.get_device_name(0))
#返回gpu名字，设备索引默认从0开始；
print(torch.cuda.current_device())
#返回当前设备索引；
device = torch.device('cuda')
#将数据转移到GPU
device = torch.device('cpu')
#将数据转移的cpu
torch.device('cuda') 与 torch.device('cuda:0')在进行计算时，对于单卡计算机而言，没有任何区别，都是唯一的那一张GPU。其中0表示GPU的索引，表示第几个GPU，在单卡机，只能是torch.device('cuda:0')，如果0换成其他数字则会报错越界。
模型可视化
from torchinfo import summary summary(model)
导入torchinfo库中的summary之后就可以查看模型的架构了
depth参数
其中summary(model)中有个参数depth，其默认为3，表示最多展示3层嵌套结果。
但是看以下transformer模型中，depth=3没法看到每个EncoderLayer中的架构
所以我们设置depth=4
更改后我们发现每个EncoderLayer中的结构也清晰的展示了，所以这个depth通过可视化的最大深度使得模型的可视化展示更加的灵活。
模型保存
在 Pytorch 中一种模型保存和加载的方式如下:
注意通常再pytorch中不是将模型的整个架构以及参数保存下来，它只是保存模型中所有的参数，模型的参数都在model.state_dict()中，将其保存到path中，通常文件后缀都以pt结尾。
所以在导入模型时实际上只是导入模型的参数，要对模型先进行实例化后才能将参数导入（换句话说得先把壳定义好才能把东西装进去），并且要满足导入的模型的架构和之前的保存时的模型架构要完全一致。
# save
torch.save(model.state_dict(), PATH)
# path文件最后的后缀维.pt
# 如torch.save(model.state_dict(),'saved_dict/basis_bert.pt')
# load
# 注意一定要先实例化模型才能将参数导入
model = MyModel(*args, **kwargs)
model.load_state_dict(torch.load(PATH))
# 如model.load_state_dict(torch.load('saved_dict/basis_bert.pt'))
model.eval()
pytorch中最常见的模型保存使用.pt或者是.pth作为模型文件扩展名。还有其他的保存数据的格式为.t7或者.pkl格式。t7文件是沿用torch7中读取模型权重的方式，而pth文件是python中存储文件的常用格式，而在keras中则是使用.h5文件 。
可以看到模型保存的是 model.state_dict() 的返回对象。 model.state_dict() 的返回对象是一个 OrderDict ，它以键值对的形式包含模型中需要保存下来的参数，例如:
class MyModule(nn.Module): def __init__(self, input_size, output_size): super(MyModule, self).__init__() self.lin = nn.Linear(input_size, output_size) def forward(self, x): return self.lin(x) module = MyModule(4, 2) print(module.state_dict())
输出结果:
模型中的参数就是线性层的 weight 和 bias.
Parameter 和 buffer
If you have parameters in your model, which should be saved and restored in the state_dict, but not trained by the optimizer, you should register them as buffers.Buffers won’t be returned in model.parameters(), so that the optimizer won’t have a change to update them.
模型中需要保存下来的参数包括两种:
-
一种是反向传播需要被optimizer更新的，称之为 parameter
-
一种是反向传播不需要被optimizer更新，称之为 buffer
第一种参数我们可以通过 model.parameters() 返回；第二种参数我们可以通过 model.buffers() 返回。因为我们的模型保存的是 state_dict 返回的 OrderDict，所以这两种参数不仅要满足是否需要被更新的要求，还需要被保存到OrderDict。
那么现在的问题是这两种参数如何创建呢，创建好了如何保存到OrderDict呢？
第一种参数有两种方式:
-
我们可以直接将模型的成员变量 通过
nn.Parameter()创建，会自动注册到parameters中，可以通过model.parameters()返回，并且这样创建的参数会自动保存到OrderDict中去； -
通过
nn.Parameter()创建普通Parameter对象，不作为模型的成员变量，然后将Parameter对象通过register_parameter()进行注册，可以通model.parameters()返回，注册后的参数也会自动保存到OrderDict中去;
第二种参数我们需要创建tensor, 然后将tensor通过register_buffer()进行注册，可以通model.buffers() 返回，注册完后参数也会自动保存到OrderDict中去。
分析一个例子：
class MyModel(nn.Module):
def __init__(self):
super(MyModel, self).__init__()
buffer = torch.randn(2, 3) # tensor
self.register_buffer('my_buffer', buffer)
self.param = nn.Parameter(torch.randn(3, 3)) # 模型的成员变量
def forward(self, x):
# 可以通过 self.param 和 self.my_buffer 访问
pass
model = MyModel()
for param in model.parameters():
print(param)
print("----------------")
for buffer in model.buffers():
print(buffer)
print("----------------")
print(model.state_dict())
输出结果:
class MyModel(nn.Module):
def __init__(self):
super(MyModel, self).__init__()
buffer = torch.randn(2, 3) # tensor
param = nn.Parameter(torch.randn(3, 3)) # 普通 Parameter 对象
self.register_buffer('my_buffer', buffer)
self.register_parameter("param", param)
def forward(self, x):
# 可以通过 self.param 和 self.my_buffer 访问
pass
model = MyModel()
for param in model.parameters():
print(param)
print("----------------")
for buffer in model.buffers():
print(buffer)
print("----------------")
print(model.state_dict())
输出:
register_buffer的函数原型:
register_buffer(name, tensor) name: string tensor: Tensor
register_parameter的函数原型:
register_parameter(name, param) name: string param: Parameter
创建第一种参数Parameter 的这两种方式有什么区别呢?
Both approaches work the same regarding training etc. There are some differences in the function calls however. Using register_parameter you have to pass the name as a string, which can make the creation of a range of parameters convenient. Besides that I think it's just coding style which one you prefer.
到这里我有两个疑问：
疑问1：为什么不把参数都设置为nn.Parameter类型，只是把不需要更新参数的设置 requires_grad=False?
疑问2：为什么不直接将不需要进行参数修改的变量作为模型类的成员变量就好了，还要进行注册?
疑问1一个解释：
两种方式都可以，使用 buffer 相当于明确告诉别人这个参数是不通过梯度下降更新的，更加可读。因为训练过程中你可以切换 requires_grad 的状态（比如说有个参数你想训练一会儿、冻结一会儿交替进行），而 buffer 永远都是不做梯度下降的。
疑问2有两个原因:
-
不进行注册，参数不能保存到
OrderDict，也就无法进行保存 -
模型进行参数在CPU和GPU移动时, 执行
model.to(device)，注册后的参数也会自动进行设备移动
看一个例子:
class MyModel(nn.Module):
def __init__(self):
super(MyModel, self).__init__()
self.my_tensor = torch.randn(1) # 参数直接作为模型类成员变量
self.register_buffer('my_buffer', torch.randn(1)) # 参数注册为 buffer
self.my_param = nn.Parameter(torch.randn(1))
def forward(self, x):
return x
model = MyModel()
print(model.state_dict())
model.cuda()
print(model.my_tensor)
print(model.my_buffer)
输出结果:
可以看到模型类的成员变量不在OrderDict中，不能进行保存；模型在进行设备移动时，模型类的成员变量没有进行移动，而加入buffer后可以进入OrderDict中，后续可以保存。
实际应用
以 transformer 中 的embedding layer 作为例子进行展示，在 "Attention is all you need" 论文中，作者对position embedding使用了固定设置：
class Embeddings(nn.Module):
def __init__(self, vocab_size, d_model, dropout=0.1, max_len=5000):
"""
Args:
vocab_size: 词典大小
d_model: 词向量维度
dropout: dropout比例
max_len: 输入序列的最大长度
"""
super(Embeddings, self).__init__()
self.embs = nn.Embedding(vocab_size, d_model) # word embedding， 反向传播需要更新
self.d_model = d_model
self.dropout = nn.Dropout(dropout)
# pe shape: (0, max_len, d_model)
pe = self._build_position_encoding(max_len, d_model)
self.register_buffer("pe", pe) # position encoding，反向传播不需要更新
def forward(self, x):
"""
Args:
x: (batch_size, seq_len)
Returns:
embed: (batch_size, seq_len, d_model)
"""
# word embedding + position encoding
embed = self.embs(x) * math.sqrt(self.d_model) + self.pe[:, :x.size(1)]
embed = self.dropout(embed)
return embed
def _build_position_encoding(self, max_len, d_model):
pe = torch.zeros(max_len, d_model)
position = torch.arange(max_len, dtype=torch.float).unsqueeze(1) # shape: (max_len, 1)
div_term = torch.exp(-torch.arange(0, d_model, 2, dtype=torch.float) * math.log(10000) / d_model) # shape: (1, ceil(d_model/2))
pe[:, 0::2] = torch.sin(position * div_term)
# d_model为偶数, 则sin 和 cos 列数相同
# d_model为奇数, 则cos 比 sin 少一列
pe[:, 1::2] = torch.cos(position * div_term) if d_model % 2 == 0 else torch.cos(position * div_term[:-1])
pe = pe.unsqueeze(0)
return pe
model = Embeddings(5, 3)
print(model.state_dict())
for name, para in model.named_parameters():
print(name, para)
for name, buffer in model.named_buffers():
print(name, buffer)
model.cuda()
print(model.state_dict())
输出结果:
总结
-
模型中需要进行更新的参数注册为Parameter，不需要进行更新的参数注册为buffer
-
模型保存的参数是
model.state_dict()返回的OrderDict -
模型进行设备移动时，模型中注册的参数(Parameter和buffer)会同时进行移动
参考: