classModel(torch.nn.Module):def__init__(self):super(Model, self).__init__()
self.linear1 = torch.nn.Linear(9,6)
self.linear2 = torch.nn.Linear(6,4)
self.linear3 = torch.nn.Linear(4,1)
self.sigmoid = torch.nn.Sigmoid()
self.relu = torch.nn.ReLU()defforward(self, x):
x = self.relu(self.linear1(x))
x = self.relu(self.linear2(x))#使用relu不会因为鞍点导致梯度消失
x = self.sigmoid(self.linear3(x))#在最后一层神经元,使用sigmoid可以把输出值限制在(0,1)之间,当然softmax也行。return x
开始训练
if __name__ =='__main__':#防止window因为多线程报错
model = Model()
criterion =torch.nn.BCELoss(size_average=True)#交叉熵
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)for epoch inrange(100):for i, data inenumerate(train_loader,0):input, label = data
y_pred = model(input)
loss = criterion(y_pred, label)print(epoch, i, loss.item())
optimizer.zero_grad()#记得梯度归零
loss.backward()
optimizer.step()