pytorch rnn

温习一下，写着玩。

import torch
import torch.nn as nn
import numpy as np
import torch.optim as optim

class RNN(nn.Module):

    def __init__(self,input_dim , hidden_dim):
        super(RNN,self).__init__()
        self._rnn = nn.RNN(input_size = input_dim , hidden_size= hidden_dim )
        self.linear = nn.Linear(hidden_dim , 1)
        self.relu = nn.ReLU()

    def forward(self , _in):
        layer1 , h = self._rnn(_in)
        layer2 = self.relu(self.linear(self.relu(layer1)))
        return layer2

    def init_weight(self):
        nn.init.normal_(self.linear.weight.data  , 0 , np.sqrt(2 / 16))
        nn.init.uniform_(self.linear.bias, 0, 0)

def getBinDict(bit_size = 16):
    max = pow(2,bit_size)
    bin_dict = {}
    for i in range(max):
        s = '{:016b}'.format(i)
        arr = np.array(list(s))
        arr = arr.astype(int)
        bin_dict[i] = arr
    return bin_dict

binary_dim = 16
int2binary = getBinDict(binary_dim)

def getBatch( batch_size):
    x = np.random.randint(0,256,[batch_size , 2])
    x_arr = np.zeros([binary_dim , batch_size , 2 ] , dtype=int)
    y_arr = np.zeros([binary_dim,batch_size,1] , dtype=int)
    for i in range(0 , binary_dim):
        batch_x_arr = np.zeros([batch_size,2] , dtype=int)
        batch_y_arr = np.zeros([batch_size,1] , dtype=int)
        for j in range(len(x)):
            batch_x_arr[j] =[int2binary[int(x[j][0])][i] , int2binary[int(x[j][1])][i]]
            batch_y_arr[j] =[int2binary[ int(x[j][0]) + int(x[j][1])][i]]

        #此处要翻转，rnn处理时是从下标为0处开始，所以要把二进制的高低位翻转
        y_arr[binary_dim - i - 1] = batch_y_arr
        x_arr[binary_dim - i - 1] = batch_x_arr
    return x_arr , y_arr , x

def getInt(y , bit_size):
    arr = np.zeros([len(y[0])])
    for i in range(len(y[0])):
        for j in range(bit_size):
            arr[i] += (int(y[j][i][0]) * pow(2 , j))
    return arr

if __name__ == '__main__':
    input_size = 2
    hidden_size = 8
    batch_size = 100
    net = RNN(input_size, hidden_size)
    net.init_weight()
    print(net)
    optimizer = optim.Adam(net.parameters(), lr=0.01, weight_decay=1e-4)
    loss_function = nn.MSELoss()#.CrossEntropyLoss()
    for i in range(100000):
        net.zero_grad()
        x ,y , t = getBatch(batch_size)
        in_x = torch.Tensor(x)
        y = torch.Tensor(y)
        output = net(in_x)
        loss = loss_function(output , y)
        loss.backward()
        optimizer.step()

        if i % 100== 0:
            output2 = torch.round(output)
            result = getInt(output2,binary_dim)
            print(t , result)
            print('iterater:%d  loss:%f'%(i , loss))