二分类问题(损失函数采用交叉熵)

import time

import matplotlib.pyplot as plt

import numpy as np

import torch

import torch.nn.functional as F

x_data = torch.tensor([[1],[2],[3]],dtype=torch.float32)

y_data = torch.tensor([[0],[0],[1]],dtype=torch.float32)

class Model(torch.nn.Module):

    def __init__(self):

        super(Model, self).__init__()

        self.linear = torch.nn.Linear(1,1)

    def forward(self,x):

        y_pred = self.linear(x)

        return torch.sigmoid(y_pred)

model = Model()

criterion = torch.nn.BCELoss(size_average=False)

optim = torch.optim.SGD(model.parameters(),lr=0.01)

for epoch in range(1000):

    y_pred = model(x_data)

    loss = criterion(y_pred,y_data)

    print(epoch, loss.item())

    # time.sleep(0.1)

    optim.zero_grad()

    loss.backward()

    optim.step()

x = np.linspace(0,10,200)

x_t = torch.tensor(x,dtype=torch.float32).view((200,1))

y_t = model(x_t)

y = y_t.data.numpy()

plt.plot(x,y)

plt.show()

你可能感兴趣的:(分类)