pytorch对比学习简单代码

以下是使用PyTorch实现一个简单的线性回归模型的示例代码：

1. 准备数据

import torch
import numpy as np

# 生成随机数据
np.random.seed(42)
x = np.random.rand(100, 1)
y = 2 * x + 0.5 + np.random.randn(100, 1) * 0.1

# 转换为Tensor
x = torch.from_numpy(x).float()
y = torch.from_numpy(y).float()

2. 定义模型

import torch.nn as nn

class LinearRegression(nn.Module):
    def __init__(self):
        super(LinearRegression, self).__init__()
        self.linear = nn.Linear(1, 1)
        
    def forward(self, x):
        out = self.linear(x)
        return out

model = LinearRegression()

3. 定义损失函数和优化器

criterion = nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

4. 训练模型

num_epochs = 1000
for epoch in range(num_epochs):
    # 前向传播，计算损失
    outputs = model(x)
    loss = criterion(outputs, y)
    
    # 反向传播，更新参数
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    
    # 输出训练过程
    if (epoch+1) % 100 == 0:
        print('Epoch [{}/{}], Loss: {:.4f}'.format(epoch+1, num_epochs, loss.item()))

5. 模型预测

# 将模型设置为评估模式
model.eval()

# 预测结果
with torch.no_grad():
    pred = model(x).numpy()

# 绘制结果
import matplotlib.pyplot as plt
plt.scatter(x.numpy(), y.numpy())
plt.plot(x.numpy(), pred, 'r-', lw=3)
plt.show()

以上代码实现了一个简单的线性回归模型，可以通过修改数据和模型的超参数来进行不同的训练和预测