Skip to content

Instantly share code, notes, and snippets.

@soumith
Created Jan 4, 2018
Embed
What would you like to do?
import torch
import torch.nn as nn
from torch.autograd import Variable
import numpy as np
BATCH_SIZE = 1
INPUT_DIM = 1
OUTPUT_DIM = 1
DTYPE = np.float32
class Net(nn.Module):
def __init__(self, input_dim, hidden_dim, output_dim, hidden_layers):
super(Net, self).__init__()
self.input_dim = input_dim
self.hidden_dim = hidden_dim
self.output_dim = output_dim
self.hidden_layers = hidden_layers
self.rnn = nn.RNN(input_dim, hidden_dim, hidden_layers)
self.h2o = nn.Linear(hidden_dim, output_dim)
def forward(self, x):
h_0 = Variable(torch.zeros(self.hidden_layers, BATCH_SIZE, self.hidden_dim))
if DTYPE == np.float32:
h_0.float()
else:
h_0.double()
output, h_t = self.rnn(x, h_0)
output = self.h2o(output)
return output
def weights_init(m):
if isinstance(m, nn.RNN):
nn.init.xavier_uniform(m.weight_ih_l0.data)
nn.init.orthogonal(m.weight_hh_l0.data)
nn.init.constant(m.bias_ih_l0.data, 0)
nn.init.constant(m.bias_hh_l0.data, 0)
if isinstance(m, nn.Linear):
nn.init.xavier_uniform(m.weight.data)
nn.init.constant(m.bias.data, 0)
data = np.loadtxt('mg17.csv', delimiter=',', dtype=DTYPE)
trX = torch.from_numpy(np.expand_dims(data[:4000, [0]], axis=1))
trY = torch.from_numpy(np.expand_dims(data[:4000, [1]], axis=1))
loss_fcn = nn.MSELoss()
model = Net(INPUT_DIM, 10, OUTPUT_DIM, 1)
model.apply(weights_init)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
for e in range(1000):
model.train()
x = Variable(trX)
y = Variable(trY)
model.zero_grad()
output = model(x)
loss = loss_fcn(output, y)
loss.backward()
optimizer.step()
print("Epoch", e + 1, "TR:", loss.cpu().data.numpy()[0])
y = Variable(trY)
model.zero_grad()
output = model(x)
loss = loss_fcn(output, y)
loss.backward()
optimizer.step()
print("Epoch", e + 1, "TR:", loss.cpu().data.numpy()[0])
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment