PyTorch

pytorch-or.py

# -*- coding: utf-8 -*-

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

x_train = torch.Tensor([
        [0, 0],
        [0, 1],
        [1, 0],
        [1, 1]
    ])

y_train = torch.Tensor([
        [0],
        [1],
        [1],
        [1]
    ])


class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__();
        self.fc1 = nn.Linear(in_features=2, out_features=3)
        self.fc2 = nn.Linear(in_features=3, out_features=1)
        
    def forward(self, x):
        out = torch.sigmoid(self.fc1(x))
        out = torch.sigmoid(self.fc2(out))
        return out

model = Net()
criterion = nn.BCELoss()
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)

for epoch in range(5000):
    optimizer.zero_grad()
    
    out = model(x_train)
    loss = criterion(out, y_train)
    loss.backward()
    
    optimizer.step()

    if epoch % 1000 == 0:
        print("epoch={}, loss={}".format(epoch, loss.data))