OpenAI Request 2.0 Warmups

warmup.py

import numpy as np

import torch
import torch.nn as nn
import torch.utils.data as data
from torch.autograd import Variable

class BinaryData(data.Dataset):

    def __init__(self):
        super(BinaryData, self).__init__()
        self.X = np.random.randint(0, 2, size=(100000, 50)).astype(np.float32)
        self.Y = np.cumsum(self.X, axis=1) % 2
        self.Y = self.Y.astype(np.int64)

    def __getitem__(self, idx):
        return self.X[idx], self.Y[idx]

    def __len__(self):
        return self.X.shape[0]


class ParityCheck(nn.Module):
    def __init__(self):
        super(ParityCheck, self).__init__()
        self.batch_size = 100
        self.hidden_size = 32
        self.lstm = nn.LSTM(32, 32, 2)
        self.input_em = nn.Embedding(2, 32)
        self.output_fc = nn.Linear(32, 2)


    def forward(self, x):
        x = self.input_em(x.long()).transpose(0, 1)
        h0 = Variable(torch.zeros(2, x.size(1), 32)).cuda()
        c0 = Variable(torch.zeros(2, x.size(1), 32)).cuda()
        yt, hn = self.lstm(x, (h0, c0))
        output = self.output_fc(yt)
        return output

def train():

    train_data = BinaryData()
    train_loader = data.DataLoader(
        train_data,
        batch_size=100,
        shuffle=True,
        num_workers=1,
        pin_memory=False
        )

    model = ParityCheck().cuda()
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(model.parameters())


    for epoch in range(10):
        losses = []
        for i, (x, y) in enumerate(train_loader):
            x_var = Variable(x).cuda()
            y_var = Variable(y).cuda()
            output = model(x_var)
            loss = sum(criterion(output[j], y_var[:,j]) for j in range(50))
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            losses.append(loss.data[0])
            if i % 100 == 0:
                print(epoch, i, np.mean(losses))


        torch.save({'state_dict': model.state_dict()}, 'save_{}.pth'.format(epoch))

def test():
    model = ParityCheck().cuda()
    checkpoint = torch.load('save_1.pth')
    model.load_state_dict(checkpoint['state_dict'])
    x = np.array([[1,0,1,0],[0,1,0,1]], dtype=np.float32)
    y = np.cumsum(x, axis=1) % 2
    x_var = Variable(torch.from_numpy(x)).cuda()
    print(model(x_var))
    print(y)


test()