PyTorch入门

1. 定义模型

from torch import nn


class LogisticRegression(nn.Module):
    def __init__(self, input_dim):
        super(LogisticRegression, self).__init__()
        
        self.linear = nn.Linear(input_dim, 1)
        self.sigmoid = nn.Sigmoid()
        
    def forward(self, x):
    	"""
    	输入 x 的维度：[batch_size, input_dim]
    	"""
        x = self.linear(x)
        x = self.sigmoid(x)
        return x

2. 训练模型

import torch
import numpy as np


device = "cuda" if torch.cuda.is_available() else "cpu"

# 获取训练数据
X = [[34.62365962451697,78.0246928153624], 
     [30.2867107622687,43.89499752400101],
     [60.18259938620976,86.3855209546826],
     [79.0327360507101,75.3443764369103]]
y = [0, 0, 1, 1]

X = torch.from_numpy(np.array(X)).to(torch.float32).to(device)
y = torch.from_numpy(np.array(y)).to(torch.float32).to(device)

# 实例化模型
model = LogisticRegression(2).to(device)
model.train()

# 定义损失函数、选择优化器
criterion = nn.BCELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=1e-3, momentum=0.9)

# 获得预测结果
y_pred = model(X)
loss = criterion(y_pred.squeeze(), y)

# 清零梯度、反向传播、更新梯度
optimizer.zero_grad()
loss.backward()
optimizer.step()

print(loss.item())

3. 评估模型

model.eval()

with torch.no_grad():
    y_pred = model(X)
    print(y_pred)

4. 保存与恢复模型

import torch

# 保存模型
torch.save(model.state_dict(), "model.pth")

# 恢复模型
model = LogisticRegression().to(device)
model.load_state_dict(torch.load("model.pth"))

5. 加载数据集

from torch.utils.data import Dataset, DataLoader
import torch
import numpy as np


class CustomDataset(Dataset):
    def __init__(self):
        X = [[34.62365962451697,78.0246928153624], 
             [30.2867107622687,43.89499752400101],
             [60.18259938620976,86.3855209546826],
             [79.0327360507101,75.3443764369103]]
        y = [0, 0, 1, 1]

        self.X = torch.from_numpy(np.array(X)).to(torch.float32).to(device)
        self.y = torch.from_numpy(np.array(y)).to(torch.float32).to(device)
    
    def __len__(self):
        return len(self.X)
    
    def __getitem__(self, index):
        x = self.X[index]
        y = self.y[index]
        return x, y

dataset = CustomDataset()
data_loader = DataLoader(dataset, batch_size=2, shuffle=True)

for X, y in data_loader:
    print(X, y)

tensor([[79.0327, 75.3444],
        [30.2867, 43.8950]]) tensor([1., 0.])
tensor([[34.6237, 78.0247],
        [60.1826, 86.3855]]) tensor([0., 1.])

从数据加载器中获取训练数据，并训练模型：

n_epoches = 100

for _ in range(n_epoches):
    for X, y in data_loader:
        y_pred = model(X)
        loss = criterion(y_pred.squeeze(), y)

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

    print(loss.item())

6. 自定义损失函数

class MyHingeLoss(torch.nn.Module):
    def __init__(self):
        super(MyHingeLoss, self).__init__()

    def forward(self, output, target):
        hinge_loss = 1 - torch.mul(output, target)
        hinge_loss[hinge_loss < 0] = 0
        return torch.mean(hinge_loss)


criterion = MyHingeLoss()
loss = criterion(torch.randn(2, 3), torch.randn(2, 3))
print(loss)