使用Pytorch构建自定义层并在模型中使用

使用Pytorch构建自定义层并在模型中使用

继承自nn.Module类，自定义名称为NoisyLinear的线性层，并在新模型定义过程中使用该自定义层。完整代码可以在jupyter nbviewer中在线访问。

import torch
import torch.nn as nn
from torch.utils.data import TensorDataset, DataLoader

import numpy as np
import matplotlib.pyplot as plt
%matplotlib inline
from mlxtend.plotting import plot_decision_regions

print(torch.__version__)
print(np.__version__)

2.0.1+cu118
1.24.4

创建一个包含有噪声的线性层

class NoisyLinear(nn.Module):
    def __init__(self, input_size, output_size, noise_stddev=0.1):
        super().__init__()
        w = torch.Tensor(input_size, output_size)
        self.w = nn.Parameter(w)
        nn.init.xavier_uniform_(self.w)
        b = torch.Tensor(output_size).fill_(0)
        self.b = nn.Parameter(b)
        self.noise_stddev = noise_stddev

    def forward(self, x, training=False):
        if training:
            noise = torch.normal(0.0, self.noise_stddev, x.shape)
            x_new = torch.add(x, noise)
        else:
            x_new = x
        return torch.add(torch.mm(x_new, self.w), self.b)

这段代码定义了一个名为 NoisyLinear 的类，它继承自 nn.Module，表示一个包含噪声的线性层。

class NoisyLinear(nn.Module):

定义一个名为 NoisyLinear 的类，它继承自 PyTorch 的 nn.Module 类。这意味着它可以被用作一种神经网络层。

    def __init__(self, input_size, output_size, noise_stddev=0.1):

初始化方法 __init__ 接受三个参数：输入大小 input_size，输出大小 output_size，以及噪声的标准差 noise_stddev（默认值为 0.1）。

        super().__init__()

调用父类 nn.Module 的初始化方法，以确保父类的相关属性和方法被正确初始化。

        w = torch.Tensor(input_size, output_size)

创建一个形状为 (input_size, output_size) 的张量 w，用于存储权重。

        self.w = nn.Parameter(w)

将权重 w 包装为 nn.Parameter，这意味着在训练过程中，PyTorch 会自动将其视为可学习参数。

        nn.init.xavier_uniform_(self.w)

使用 Xavier 均匀分布对权重 self.w 进行初始化。这是一种常用的初始化方法，有助于保持神经网络中信号的方差。

        b = torch.Tensor(output_size).fill_(0)

创建一个形状为 (output_size,) 的张量 b，并将其填充为 0，用于存储偏置。

        self.b = nn.Parameter(b)

将偏置 b 包装为 nn.Parameter，使其在训练过程中也是可学习的。

        self.noise_stddev = noise_stddev

将噪声的标准差 noise_stddev 存储为类的一个属性，用于后续的噪声计算。

    def forward(self, x, training=False):

定义前向传播方法 forward，接受输入 x 和一个布尔参数 training，指示当前是否在训练模式下。

        if training:

检查当前是否处于训练模式。

            noise = torch.normal(0.0, self.noise_stddev, x.shape)

如果是训练模式，则创建一个与输入 x 形状相同的噪声张量 noise，其服从均值为 0、标准差为 self.noise_stddev 的正态分布。

            x_new = torch.add(x, noise)

将噪声添加到输入 x 上，得到新的输入 x_new。

        else:

如果不是训练模式，则执行以下代码。

            x_new = x

在非训练模式下，x_new 直接设置为输入 x，即没有添加噪声。

        return torch.add(torch.mm(x_new, self.w), self.b)

计算输出：首先用 torch.mm 进行矩阵乘法（x_new 和权重 self.w），然后将偏置 self.b 添加到结果中。最后返回计算出的输出。

总结来说，这个类实现了一个带噪声的线性变换，在线性层中可以根据训练模式选择性地添加噪声。

# 上述层的使用示例.
# 1、实例化这个层，并调用三次.
torch.manual_seed(1)

noisy_layer = NoisyLinear(4, 2)
x = torch.zeros((1, 4))
print(noisy_layer(x, training=True))

print(noisy_layer(x, training=True))

print(noisy_layer(x, training=False))

tensor([[ 0.1154, -0.0598]], grad_fn=<AddBackward0>)
tensor([[ 0.0432, -0.0375]], grad_fn=<AddBackward0>)
tensor([[0., 0.]], grad_fn=<AddBackward0>)

在一个示例数据上，构建一个包含该自定义层的模型

# 生成一个示例数据.
np.random.seed(1)
torch.manual_seed(1)
x = np.random.uniform(low=-1, high=1, size=(200, 2))
y = np.ones(len(x))
y[x[:, 0] * x[:, 1]<0] = 0

n_train = 100
x_train = torch.tensor(x[:n_train, :], dtype=torch.float32)
y_train = torch.tensor(y[:n_train], dtype=torch.float32)
x_valid = torch.tensor(x[n_train:, :], dtype=torch.float32)
y_valid = torch.tensor(y[n_train:], dtype=torch.float32)

fig = plt.figure(figsize=(6, 6))
plt.plot(x[y==0, 0], 
         x[y==0, 1], 'o', alpha=0.75, markersize=10)
plt.plot(x[y==1, 0], 
         x[y==1, 1], '<', alpha=0.75, markersize=10)
plt.xlabel(r'$x_1$', size=15)
plt.ylabel(r'$x_2$', size=15)
plt.tight_layout()
plt.show()

# 创建一个DataLoader.
train_ds = TensorDataset(x_train, y_train)
batch_size = 2
torch.manual_seed(1)

# 使用DataLoader加载数据，batchsize为2.
train_dl = DataLoader(train_ds, batch_size, shuffle=True)

# 创建一个新的模型，并且调用上述的自定义层.
class MyNoiseModule(nn.Module):
    def __init__(self):
        super().__init__()
        self.l1 = NoisyLinear(2, 4, 0.07)
        self.a1 = nn.ReLU()
        self.l2 = nn.Linear(4, 4)
        self.a2 = nn.ReLU()
        self.l3 = nn.Linear(4, 1)
        self.a3 = nn.Sigmoid()

    def forward(self, x, training=False):
        x = self.l1(x, training)
        x = self.a1(x)
        x = self.l2(x)
        x = self.a2(x)
        x = self.l3(x)
        x = self.a3(x)
        return x
    
    def predict(self, x):
        self.eval()
        with torch.no_grad():
            x = torch.tensor(x, dtype=torch.float32)
            pred = self.forward(x)[:, 0]
            return (pred>=0.5).float()

# 模型实例化.
torch.manual_seed(1)
model = MyNoiseModule()
model

MyNoiseModule(
  (l1): NoisyLinear()
  (a1): ReLU()
  (l2): Linear(in_features=4, out_features=4, bias=True)
  (a2): ReLU()
  (l3): Linear(in_features=4, out_features=1, bias=True)
  (a3): Sigmoid()
)

# 3.在训练training batch上计算预测结果.
loss_fn = nn.BCELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.015)

# 模型训练，设置epochs=200
torch.manual_seed(1)
num_epochs = 200

def train(model, num_epochs, train_dl, x_valid, y_valid):
    loss_hist_train = [0] * num_epochs
    acc_hist_train = [0] * num_epochs

    loss_hist_valid = [0] * num_epochs
    acc_hist_valid = [0] * num_epochs

    for epoch in range(num_epochs):
        for x_batch, y_batch in train_dl:
            pred = model(x_batch, True)[:, 0]
            loss = loss_fn(pred, y_batch)
            loss.backward()
            optimizer.step()
            optimizer.zero_grad()

            loss_hist_train[epoch] += loss.item()
            is_correct = ((pred>=0.5).float() == y_batch).float()
            acc_hist_train[epoch] += is_correct.mean()
        loss_hist_train[epoch] /= n_train/batch_size
        acc_hist_train[epoch] /= n_train/batch_size

        pred = model(x_valid)[:, 0]
        loss = loss_fn(pred, y_valid)
        loss_hist_valid[epoch] = loss.item()
        is_correct = ((pred>=0.5).float() == y_valid).float()
        acc_hist_valid[epoch] += is_correct.mean()
    return loss_hist_train, loss_hist_valid, \
            acc_hist_train, acc_hist_valid

history = train(model, num_epochs, train_dl, x_valid, y_valid)

# 绘制决策边界.
fig = plt.figure(figsize=(16, 4))
ax = fig.add_subplot(1, 3, 1)
plt.plot(history[0], lw=4)
plt.plot(history[1], lw=4)
plt.legend(['Train loss', 'Validation loss'], fontsize=15)
ax.set_xlabel('Epochs', size=15)

ax = fig.add_subplot(1, 3, 2)
plt.plot(history[2], lw=4)
plt.plot(history[3], lw=4)
plt.legend(['Train acc.', 'Validation acc.'], fontsize=15)
ax.set_xlabel('Epochs', size=15)

ax = fig.add_subplot(1, 3, 3)
plot_decision_regions(X=x_valid.numpy(), 
                      y=y_valid.numpy().astype(np.int64),
                      clf=model)
ax.set_xlabel(r'$x_1$', size=15)
ax.xaxis.set_label_coords(1, -0.025)
ax.set_ylabel(r'$x_2$', size=15)
ax.yaxis.set_label_coords(-0.025, 1)
plt.show()