かってきままな日々(12-18[長年日記])

2019-12-18 (We)

_ PyTorch で RNN

やっとできた。ついでに GPU に載せた。

#!/usr/bin/env python

import torch
import numpy as np
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import math
import random
import matplotlib.pyplot as plt

hidden_size = 8
seqlen = 8
batch_size = 128
n_epoch = 1000

N_PIE = 128
N_ROUND = 16

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.rnn = nn.RNN(input_size=2, hidden_size=hidden_size, batch_first=True)
        self.fc1 = nn.Linear(hidden_size, 2)

    def forward(self, x):
        h0 = torch.zeros(1, x.size(0), hidden_size).to('cuda')
        output, hn = self.rnn(x, h0)
        x = self.fc1(output[:,x.size(1) - 1,:])
        return x

net = Net().to('cuda')

criterion = nn.MSELoss()
optimizer = optim.SGD(net.parameters(), lr=0.1, momentum=0.9)

def calc_point(no, noisy):
    theta = math.pi * 2 / N_PIE * no
    x = math.cos(theta)
    y = math.sin(theta)
    if noisy:
        x += random.random() * 0.4 - 0.2
        y += random.random() * 0.4 - 0.2
    return (x, y)

def make_seq(beg, size, noisy=False):
    seq = [ calc_point(no, noisy) for no in range(beg, beg + size) ]
    return (seq, calc_point(beg + size, noisy))

inputs = []
expecteds = []
for i in range(N_PIE * N_ROUND):
    seq, expected = make_seq(i, seqlen, True)
    expecteds.append(expected)
    inputs.append(seq)

inputs = torch.tensor(inputs).to('cuda')
expecteds = torch.tensor(expecteds).to('cuda')

for epoch in range(n_epoch):
    perm = torch.randperm(inputs.size(0)).to('cuda')

    beg = 0
    while beg + batch_size < inputs.size(0):
        batch = perm[beg:beg+batch_size]

        # zero the parameter gradients
        optimizer.zero_grad()

        # forward + backward + optimize
        outputs = net(inputs[batch])
        loss = criterion(outputs, expecteds[batch])
        loss.backward()
        optimizer.step()

        beg += batch_size

    print(epoch, loss.item())

print('Training finished.')


seq, _ = make_seq(0, seqlen)

xs = []
ys = []
for n in range(1024):
    outputs = net(torch.tensor([seq]).to('cuda'))
    x = outputs[0][0].item()
    y = outputs[0][1].item()
    print('%.4f %.4f' % (x, y))
    xs.append(x)
    ys.append(y)
    seq = seq[1:]
    seq.append([x, y])

plt.plot(xs, ys)
plt.savefig('test.png')

円周を N_PIE 個に分割し、seqlen 個の連続した点から次の点を予測する。古い点を捨てながら新しい点を追加しながらこれを連続してやれば、綺麗な円が描けるはず。

というプログラム。結果は↓これ。