先尝试写一个简单的线性计算，用神经网络拟合学习。

1
2
3
4
5
6
7
8
9
10
# 生成数据
import torch
# 给定权重矩阵
W = torch.tensor([[0.3, 0.4],[0.5,0.6],[0.6,0.2]])

num_samples = 100
input_dim = 3
x = torch.rand((num_samples, input_dim))

y = torch.matmul(W.t(), x.t()).t()

一会儿我们用这些数据取拟合权重矩阵 $W$

pytorch中的自定义神经网络类需要继承 nn.Module ，并且调用父类的构造函数 __init__() 来完成一些参数和方法的初始化。

最简单的网络我们只需定义网络中的层结构，如这里我加了一层线性层。以及前向传播的 forward 函数即可。

如果需要自定义损失函数等，会在之后的文章中介绍。

1
2
3
4
5
6
7
8
9
10
import torch.nn as nn
class Simple_nn(nn.Module):
    def __init__(self):
        super(Simple_nn, self).__init__()
        self.W = nn.Linear(input_dim, output_dim, bias=False)
    def forward(self, X):
        # X : [batch_size, input_dim]
        # 在forward中的X还是带有batch_size这一维度的，不过在经过链接层的时候，pytorch会自动处理批次，不用显式考虑。
        output = self.W(X)
        return output

训练需要数据，我们把数据以batch的形式送进网络进行训练。目前我们不涉及dataset，dataloader的使用。

先写一个随机取batch的函数

1
2
3
4
5
6
7
8
9
10
11
12
13
batch_size = 4 # mini-batch size

def random_batch():
    random_x = torch.zeros((batch_size, input_dim))
    random_y = torch.zeros((batch_size, output_dim))
    # 定义一个随机抽取batch
    random_index = np.random.choice(range(len(x)), batch_size, replace=False)

    for i, index in enumerate(random_index):
        random_x[i] = x[index]
        random_y[i] = y[index]
                
    return random_x, random_y

然后写一个训练函数

pytorch封装了很多操作，比如说反向传播 backward 是不需要自己写的，只需要调用 loss.backward() 即可

随后，调用 torch.optim 中被称为优化器的工具就可以更新参数

像这里的 optim.SGD 就是随机梯度下降的参数更新方法。

注意，模型直到 optimizer.step() 这步才正式更新模型参数。这两步的分离是为了提供更大的灵活性，比如可以多次调用 loss.backward()累积梯度，然后在特定时刻执行一次 optimizer.step()

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
import torch.optim as optim

# 创建模型实例
model = Simple_nn()
# 定义损失函数 
criterion = nn.MSELoss()
# 定义优化器
optimizer = optim.SGD(model.parameters(), lr=0.01)

# Training
for epoch in range(5000):
    input_batch, output_batch = random_batch()

    # 清除梯度缓存
    optimizer.zero_grad()
    # 前向传播
    output_pred = model(input_batch)
    # 计算损失函数
    loss = criterion(output_pred, output_batch)

    if (epoch + 1) % 1000 == 0:
            print('Epoch:', '%04d' % (epoch + 1), 'cost =', '{:.6f}'.format(loss))

    loss.backward()
    optimizer.step()

print(model.W.weight.t())

最后得到输出，和我们一开始定义的 $W$ 是很接近的：

1
2
3
tensor([[0.2985, 0.4066],
        [0.5058, 0.5921],
        [0.5955, 0.2015]], grad_fn=<TBackward0>)