PyTorch实现二分类任务

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

'''定义模拟数据集'''
inputs = torch.randn(5, 10)  # 定义5个样本,每个样本的特征数位10
labels = torch.tensor([0, 1, 0, 1, 1])  # 每个样本的标签,二分类,所以为0或1

'''定义模型'''  
class SimpleModel(nn.Module):
    def __init__(self, input_size, output_size=1):
        super().__init__()
        self.linear = nn.Linear(input_size, output_size)    
        self.sigmoid = nn.Sigmoid() 
        
    def forward(self, x):
        x = self.linear(x)
        x = self.sigmoid(x)
        return x

'''实例化模型'''
model = SimpleModel(input_size=10, output_size=1)

'''定义二元交叉熵损失函数'''
criterion = nn.BCELoss()

'''定义优化器'''    
optimizer = optim.SGD(model.parameters(), lr=0.01)

'''训练'''
epoch = 50000
for i in range(epoch):
    '''前向传播'''
    # 执行过程中遇到shape报错:纠正后加了一个reshape(-1)    
    y = model(inputs).reshape(-1)    
    
    '''计算损失'''
    # 执行过程中遇到数据类型报错,因此加了一个.float()
    loss = criterion(y.float(), labels.float())
    print(f'Epoch: {i + 1}/{epoch}, Loss: {loss.item()}')   
    
    '''反向传播 和 更新参数'''  
    optimizer.zero_grad()  # 梯度清零
    loss.backward()  # 计算梯度 
    optimizer.step()  # 更新参数

  

posted @ 2024-06-07 17:29  映辉  阅读(15)  评论(0编辑  收藏  举报