pytorch基础知识学习（1）

概要

本文主要是记录 pytorch 的基础学习，参考文章为：一文理解PyTorch:附代码实例

1、主程序：hewu_train.py

"""
export CUDA_VISIBLE_DEVICES='7'
cd "$(find ~/ -type f -name "hewu_train.py" -exec dirname {} \; -quit)"
/ssddata/home/hewu/anaconda3/envs/torch==1.9.0+cu102/bin/python hewu_train.py

"""

import torch
import torch.optim as optim
import torch.nn as nn

import hewu_input_data
import hewu_models



# 加载输入数据，并指定设备
device = "cuda" if torch.cuda.is_available() else "cpu"
print("device is : ",device)

train_loader = hewu_input_data.train_loader
val_loader = hewu_input_data.val_loader


# 创建模型
torch.manual_seed(42)
if False:
    model = hewu_models.LayerLinearRegression().to(device)
else:
    # 用顺序模型代替嵌套模型
    model = nn.Sequential(nn.Linear(1, 1)).to(device)
print(model.state_dict())


lr = 1e-1
n_epochs = 1000

loss_fn = nn.MSELoss(reduction='mean')
optimizer = optim.SGD(model.parameters(), lr=lr)


# Creates the train_step function for our model, loss function and optimizer
losses = []
val_losses = []
train_step = hewu_models.make_train_step(model, loss_fn, optimizer)

for epoch in range(n_epochs):
    for x_batch, y_batch in train_loader:
        # 整个数据以及小批量数据集都在CPU中，因此我们需要把将要使用的小批量数据集送到GPU中
        x_batch = x_batch.to(device)
        y_batch = y_batch.to(device)

        loss = train_step(x_batch, y_batch) # 使用小批量来计算损失，更新参数
        losses.append(loss)

    if epoch % (n_epochs/5) == 0:    # 只验证5次
        with torch.no_grad():
            for x_val, y_val in val_loader:
                x_val = x_val.to(device)
                y_val = y_val.to(device)

                model.eval()

                yhat = model(x_val)
                val_loss = loss_fn(y_val, yhat)
                val_losses.append(val_loss.item())


# Checks model's parameters
print(model.state_dict())
print(losses[0:n_epochs:(n_epochs//5)])
print(val_losses)

2、模型：hewu_models.py

import torch.nn as nn


class LayerLinearRegression(nn.Module):
    def __init__(self):
        super().__init__()  # 调用父类构造函数，初始化父类属性，确保子类能继承并使用父类属性
        self.linear = nn.Linear(1, 1)

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


def make_train_step(model, loss_fn, optimizer):
    # Builds function that performs a step in the train loop
    def train_step(x, y):   # 嵌套函数可以使用外层函数的局部变量和参数，使用嵌套函数是为了模块化
        # Sets model to TRAIN mode
        model.train()  # 唯一作用是将模型设置为训练模式，因为模型如果有 dropout 机制的话，训练和评估是不同的
        # Makes predictions
        yhat = model(x)
        """
        model(x)等价于model.forward(x)，这是因为在父类 nn.Module 中使用了 __call__ 函数
        详情请参考：https://zhuanlan.zhihu.com/p/357021687
                  https://zhuanlan.zhihu.com/p/356059224
        """
        # Computes loss
        loss = loss_fn(y, yhat)
        # Computes gradients
        loss.backward()
        # Updates parameters and zeroes gradients
        optimizer.step()
        optimizer.zero_grad()
        # Returns the loss
        return loss.item()  # item()是PyTorch的一个方法，用于从张量中提取单一值，并将其转换为Python浮点数。

    # Returns the function that will be called inside the train loop
    """
    这种技术非常有用：
    在一个函数内部定义了另一个函数，实现函数嵌套，并且将这个内部函数作为结果返回给外部。这意味着外部函数实际上生成了一个新的函数，
    并且这个新函数可以在外部函数之外的地方调用和使用。
    此代码中，make_train_step() 是一个外部函数，它接受一个参数 (model, loss_fn, optimizer)。
    在内部，它定义了一个内部函数 train_step(x, y)，并将 train_step 作为结果返回给外部。
    这就意味着在调用 make_train_step() 后，我们实际上获得了一个新的函数 train_step，该函数可以在程序的其他地方调用，就像任何其他函数一样。
    """
    return train_step

3、数据集：hewu_input_data.py

import numpy as np
import torch
from torch.utils.data import Dataset, TensorDataset
from torch.utils.data.dataset import random_split
from torch.utils.data import DataLoader
import sys

class CustomDataset(Dataset):
    def __init__(self, x_tensor, y_tensor):
        self.x = x_tensor
        self.y = y_tensor

    def __getitem__(self, index):
        return (self.x[index], self.y[index])

    def __len__(self):
        return len(self.x)



# 数据生成
# Data Generation
np.random.seed(42)
x = np.random.rand(100, 1)
y = 1 + 2 * x

# 为什么放 CPU ,不把全部训练数据加载到 GPU 中去，是为了把 GPU 的内存空出来
x_tensor = torch.from_numpy(x).float()
y_tensor = torch.from_numpy(y).float()

if True:   # 自定义数据集
    dataset = CustomDataset(x_tensor, y_tensor)
else:      # 张量比较少
    dataset = TensorDataset(x_tensor, y_tensor)

train_dataset, val_dataset = random_split(dataset, [80, 20])


train_loader = DataLoader(dataset=train_dataset, batch_size=16)
val_loader = DataLoader(dataset=val_dataset, batch_size=20)