Torch 基础

发表于 更新于 分类于 本文字数： 853 阅读时长 ≈ 3 分钟
本文介绍了 PyTorch 的基础知识，包括输入输出、构建网络结构、优化器与损失函数、保存与加载模型、训练一个简单的网络学习 XOR。

Troch 输入输出

创建一个 tensor 变量

1
x = torch.tensor([1, 2, 3])

转换类型

1
x = x.float()

并使用 cuda() 函数将变量转移到 GPU 上。

1
x = x.cuda()

同样，我们可以使用 cpu() 函数将变量转移到 CPU 上。

1
x = x.cpu()

完整的代码如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
import torch

x = torch.tensor([1, 2, 3, 4])
y = torch.tensor([[1, 2], [3, 4]])

print("Before cuda")
print(f'x: {x}, {x.type()}, {x.size()}')
print(f'y: {y}, {y.type()}, {y.size()}')
print("#######################################################")

x = x.float()
y = y.float()

x = x.cuda()
y = y.cuda()

print("After cuda")
print(f'x: {x}, {x.type()}, {x.size()}')
print(f'y: {y}, {y.type()}, {y.size()}')

输出结果如下：

1
2
3
4
5
6
7
8
9
Before cuda
x: tensor([1, 2, 3, 4]), torch.LongTensor, torch.Size([4])
y: tensor([[1, 2],
        [3, 4]]), torch.LongTensor, torch.Size([2, 2])
#######################################################
After cuda
x: tensor([1., 2., 3., 4.], device='cuda:0'), torch.cuda.FloatTensor, torch.Size([4])
y: tensor([[1., 2.],
        [3., 4.]], device='cuda:0'), torch.cuda.FloatTensor, torch.Size([2, 2])

构建网络结构

使用 torch.nn 构建网络结构。

1
2
3
4
5
6
7
8
import torch.nn as nn

myNet = nn.Sequential(
    nn.Linear(2, 3),
    nn.ReLU(),
    nn.Linear(3, 1),
    nn.Sigmoid()
)

使用类的形式构建网络结构。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
class myNet(nn.Module):
    def __init__(self):
        super(myNet, self).__init__()
        self.fc1 = nn.Linear(2, 3)
        self.fc2 = nn.Linear(3, 1)
        self.relu = nn.ReLU()
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        x = self.sigmoid(x)
        return x

上面两种方式都可以，但是第二种方式更加灵活，可以自定义网络结构。但是，它们都使用 CPU 计算，如果需要使用 GPU 计算，需要在构建网络结构的时候加上 .cuda()。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
class myNet(nn.Module):
    def __init__(self):
        super(myNet, self).__init__()
        self.fc1 = nn.Linear(2, 3).cuda()
        self.fc2 = nn.Linear(3, 1).cuda()
        self.relu = nn.ReLU().cuda()
        self.sigmoid = nn.Sigmoid().cuda()

    def forward(self, x):
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        x = self.sigmoid(x)
        return x

优化器与损失函数

使用 torch.optim 构建优化器。

这里使用了随机梯度下降法（SGD）。还有其他的优化器，如 Adam、RMSProp 等。参考：torch.optim

1
2
3
import torch.optim as optim

optimizer = optim.SGD(myNet.parameters(), lr=0.01)

使用 torch.nn 构建损失函数。

1
2
3
import torch.nn as nn

criterion = nn.MSELoss()

保存与加载模型

保存模型

1
2
3
4
5
6
# 这种方式只保存模型的参数
torch.save(myNet.state_dict(), 'myNet.pth')

# or
# 这种方式保存整个模型
torch.save(myNet, 'myNet.pth')

加载模型

1
myNet.load_state_dict(torch.load('myNet.pth'))

训练一个简单的网络学习 XOR

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
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
import torch
import torch.nn as nn


class MyNet(nn.Module):
    def __init__(self):
        super(MyNet, self).__init__()
        self.net = nn.Sequential(
            nn.Linear(2, 20),
            nn.ReLU(),
            nn.Linear(20, 1),
            nn.Sigmoid()
        )
        self.optimizer = torch.optim.SGD(self.net.parameters(), lr=0.05)
        self.loss_func = nn.MSELoss()

    def forward(self, x):
        return self.net(x)

    def train_xor(self, x, y, epochs=5000):
        for epoch in range(epochs):
            out = self.forward(x)
            loss = self.loss_func(out, y)

            self.optimizer.zero_grad()
            loss.backward()
            self.optimizer.step()

            if epoch % 1000 == 0:
                print(f'epoch: {epoch}, loss: {loss}')

        out = self.forward(x)
        print(f'out: {out.data}')

        torch.save(self.net, 'net_EOR_CPU.pkl')


def main():
    x = [[0, 0], [0, 1], [1, 0], [1, 1]]
    y = [[0], [1], [1], [0]]
    x_tensor = torch.tensor(x).float()
    y_tensor = torch.tensor(y).float()

    net = MyNet()
    print(net)

    net.train_xor(x_tensor, y_tensor)


if __name__ == '__main__':
    main()

测试模型

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
30
31
32
import torch


class MyNet:
    def __init__(self, model_path):
        self.net = torch.load(model_path, map_location='cpu')

    def predict(self, x):
        out = self.net(x)
        outfinal = out.data
        return outfinal


def main():
    x = [[0, 0], [0, 1], [1, 0], [1, 1]]
    x_tensor = torch.tensor(x).float()

    net = MyNet('net_EOR_CPU.pkl')
    outfinal = net.predict(x_tensor)

    print(f'out: {outfinal}')

    for i in range(4):
        if outfinal[i] < 0.5:
            print(f'out: {outfinal[i]}, 0')
        else:
            print(f'out: {outfinal[i]}, 1')


if __name__ == '__main__':
    main()

这个例子中，我们使用了 torch.nn 构建了一个网络结构，使用了 torch.optim 构建了一个优化器，使用了 torch.nn 构建了一个损失函数，使用了 torch.save 保存了模型，使用了 torch.load 加载了模型。

并且使用的是 CPU 计算，如果需要使用 GPU 计算，需要在构建网络结构的时候加上 .cuda()。

最后 print 时，将结果cpu化，即 print(net(x).cpu())。这样才能在 CPU 上打印出结果。

参考

  • PyTorch官方文档

  • PyTorch中文文档

  • ⟪Python 神经网络入门与实战⟫ ———— 王凯编著，北京大学出版社，2020年11月第1版