PyTorch 使用心得

时间:2021-03-20 22:34:23

PyTorch 使用心得

模板

import torch.nn as nn
import torch.optim as optim class Model(nn.Module):
def __init__(self):
super(Model, self).__init__()
# self.hidden_layer = nn.Linear(in_dim, out_dim) def forward(self, x):
# out = self.hidden_layer(x)
return out model = Model()
optimizer = optim.Adam(model.parameters(), lr=lr)
criticism = nn.MSELoss() for i in range(epoch):
optimizer.zero_grad() y_pred = model(x)
loss = criticism(y_pred, y) loss.backward()
optimizer.step() # test

标准库

计算设备(张量类型)

设置 tensor 的默认类型

torch.set_default_tensor_type('torch.cuda.FloatTensor')  # 默认使用GPU

设置程序优先使用 CUDA (e.g. GPU)

gpu = torch.device("cuda" if torch.cuda.is_available() else "cpu")

Cosine作为损失函数

>>> x = torch.randn(1, 100)
>>> y = torch.randn(1, 100)
>>> torch.nn.functional.cosine_similarity(x,y)
tensor([-0.1123])

变量

类型转换

Tensor to float

a = torch.Tensor(3.14)  # 'a' is a type of Tensor
b = float(a) # b is 3.14, a type of float