模型 ¶

# 基本配置
## 常用包
import torch

## 超参数配置，有时也写入特定的文件中
batch_size = 16
lr = 1e-4 # 学习率
max_epoch = 100 # 训练轮数

## 硬件配置
def get_device(k=2):
    """
    智能选择训练设备
    - 如果有多个GPU，使用最后 k 个
    - 如果只有一个GPU，使用这一个
    - 如果没有GPU，回退到CPU
    """
    if torch.cuda.is_available():
        gpu_count = torch.cuda.device_count()
        print(f"GPU数量: {gpu_count}")
        if gpu_count >= k:
            device_ids = list(range(gpu_count-k, gpu_count))
        else:
            device_ids = list(range(gpu_count))
        print(f"使用GPU: {device_ids}")
        return torch.device(f"cuda:{device_ids[-1]}"), device_ids
    else:
        # 回退到CPU
        print("使用CPU")
        return torch.device("cpu"), None

device, device_ids = get_device(2)

# 简单网络构建
import torch
from torch import nn

class MLP(nn.Module):
    """
    Multi Layer Perceptron
    """
    def __init__(self, input_dim=784, hidden_dim=256, output_dim=10, **kwargs):
        super(MLP, self).__init__(**kwargs) # 继承父类的必要初始化
        self.hidden = nn.Linear(input_dim, hidden_dim) # 全连接层
        self.act = nn.ReLU() # 激活层
        self.output = nn.Linear(hidden_dim, output_dim)

    def forward(self, x):
        """
        A classic forward function, as neuron
        """
        o = self.act(self.hidden(x))
        return self.output(o)

X = torch.rand(2, 784)
mlp = MLP(input_dim=784, hidden_dim=256, output_dim=10)
mlp, mlp(X)

# 自定义网络层
class MCL(nn.Module):
    """
    Mean Center Layer
    """
    def __init__(self, **kwargs):
        super(MCL, self).__init__(**kwargs)
    def forward(self, x):
        return x - x.mean()
    
X = torch.arange(5, dtype=torch.float)
mcl = MCL()
mcl, mcl(X)

# 二维卷积层（含可学习参数）
# 注意下面实际操作使用了 nn.Conv2d，而不是我们的这个最小实现
class MyConv2d(nn.Module):
    """
    A simple convolutional layer
    :param kernel_size: kernel size
    """
    def __init__(self, kernel_size):
        super(MyConv2d, self).__init__()
        # random generate initial parameters
        self.weight = nn.Parameter(torch.randn(kernel_size))
        self.bias = nn.Parameter(torch.randn(1))

    def forward(self, x):
        return self.corr2d(x, self.weight) + self.bias
    
    def corr2d(X, K):
        """
        :param X: original matrix
        :param K: correlation kernel
        """
        h, w = K.shape
        X, K = X.float(), K.float()
        Y = torch.zeros((X.shape[0] - h + 1, X.shape[1] - w + 1))
        for i in range(Y.shape[0]):
            for j in range(Y.shape[1]):
                Y[i, j] = (X[i: i + h, j: j + w] * K).sum()
        return Y

# 定义一个函数来计算卷积层。它对输入和输出做相应的升维和降维
def comp_conv2d(conv2d, X):
    X = X.view((1, 1) + X.shape) # (1, 1)代表批量大小和通道数
    Y = conv2d(X)
    return Y.view(Y.shape[2:]) # 排除不关心的前两维:批量和通道

X = torch.rand(8, 8)
# 注意这里是两侧分别填充1⾏或列，所以在两侧一共填充2⾏或列
conv2d = nn.Conv2d(in_channels=1, out_channels=1, kernel_size=3,padding=1)
print(comp_conv2d(conv2d, X).shape)
# 当卷积核的高和宽不一致时填充也可以改变
conv2d = nn.Conv2d(in_channels=1, out_channels=1, kernel_size=(5, 3), padding=(2, 1))
print(comp_conv2d(conv2d, X).shape)
# 利用 stride 改变每次采样时滑动的步长
conv2d = nn.Conv2d(1, 1, kernel_size=(3, 5), padding=(0, 1), stride=(3, 4))
print(comp_conv2d(conv2d, X).shape)

# 池化层
# 池化层可以看作是卷积核特殊的卷积层，但是进行的操作一般是顺序无关的操作
class MyPool2d(nn.Module):
    """
    simple pooling layer
    """
    def __init__(self, pool_size, mode="mean"):
        super(MyPool2d, self).__init__()
        self.pool_size = pool_size
        self.mode = mode
    
    def forward(self, x):
        return self.pool2d(x, self.pool_size, mode=self.mode)

    def pool2d(self, X, pool_size, mode="mean"):
        p_h, p_w = pool_size
        Y = torch.zeros((X.shape[0] - p_h + 1, X.shape[1] - p_w + 1))
        for i in range(Y.shape[0]):
            for j in range(Y.shape[1]):
                if mode == 'mean':
                    Y[i, j] = X[i: i + p_h, j: j + p_w].mean()
                elif mode == 'max':
                    Y[i, j] = X[i: i + p_h, j: j + p_w].max()
        return Y

X = torch.tensor([[0, 1, 2], [3, 4, 5], [6, 7, 8]], dtype=torch.float)
mypool2d = MyPool2d(pool_size=(2, 2))
mypool2d, mypool2d(X)

# 模型初始化
# 在 https://pytorch.org/docs/stable/nn.init.html 可以看到模型初始化的一些方法；以 `_` 结尾表明这些方法对 tensor 本身进行操作，会自动选择合适的方法进行初始化；当然要主动选择初始化方式也没有问题。
conv = nn.Conv2d(1, 3, 3)
linear = nn.Linear(10, 1)
conv.weight.data # 默认

# LeNet
import torch.nn.functional as F

class LeNet(nn.Module):

    def __init__(self):
        super(LeNet, self).__init__()
        # 输入图像channel：1；输出channel：6；5x5卷积核
        self.conv1 = nn.Conv2d(1, 6, 5)
        self.conv2 = nn.Conv2d(6, 16, 5)
        # an affine operation: y = Wx + b
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        # 2x2 Max pooling
        x = F.max_pool2d(F.relu(self.conv1(x)), (2, 2))
        # 如果是方阵,则可以只使用一个数字进行定义
        x = F.max_pool2d(F.relu(self.conv2(x)), 2)
        x = x.view(-1, self.num_flat_features(x))
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

    def num_flat_features(self, x):
        size = x.size()[1:]  # 除去批处理维度的其他所有维度
        num_features = 1
        for s in size:
            num_features *= s
        return num_features
    
lenet = LeNet()
params = list(lenet.parameters())
# Net 自动将 Conv2d 和 Linear 的权重和偏置注入到模型中，5层网络，len(params)=10
len(params), params[0].size(), params[1].size()
# (10, torch.Size([6, 1, 5, 5]), torch.Size([6]))
output = lenet(torch.randn(1, 1, 32, 32))
lenet.zero_grad()
output.backward(torch.randn(1, 10))