Replace tensorflow with pytorch

This commit is contained in:
Shangtong Zhang
2017-05-04 14:05:40 -06:00
parent 3f13ebc711
commit 966e3ebd60
5 changed files with 88 additions and 203 deletions
+1
View File
@@ -7,6 +7,7 @@ __pycache__/
exp_*
upload.py
*.sh
data
# C extensions
*.so
-115
View File
@@ -1,115 +0,0 @@
#######################################################################
# Copyright (C) 2017 Shangtong Zhang(zhangshangtong.cpp@gmail.com) #
# Permission given to modify the code as long as you keep this #
# declaration at the top #
#######################################################################
import tensorflow as tf
from common import *
import numpy as np
class CrossProp:
def __init__(self, name, dim_in, dim_hidden, dim_out, learning_rate, gate=Relu(),
initializer=tf.random_normal_initializer(), bottom_layer=None,
optimizer=None, output_layer='MSE'):
self.learning_rate = learning_rate
self.lam = 0
self.output_layer = output_layer
if optimizer is None:
optimizer = tf.train.GradientDescentOptimizer(learning_rate=self.learning_rate)
if bottom_layer is None:
self.x = tf.placeholder(tf.float32, shape=(None, dim_in))
var_in = self.x
trainable_vars = []
else:
self.x = bottom_layer.x
var_in = bottom_layer.var_out
trainable_vars = bottom_layer.trainable_vars
self.bottom_layer = bottom_layer
self.h = tf.placeholder(tf.float32, shape=(dim_hidden, dim_out))
self.target = tf.placeholder(tf.float32, shape=(None, dim_out))
U, b_hidden, net, phi, W, b_out, y =\
crossprop_layer(name, 'crossprop_layer', var_in, dim_in, dim_hidden, dim_out, gate.gate_fun, initializer)
if self.output_layer == 'CE':
self.pred = tf.nn.softmax(y)
ce_loss = tf.nn.softmax_cross_entropy_with_logits(logits=y, labels=self.target)
self.loss = tf.reduce_mean(ce_loss)
self.total_loss = tf.reduce_sum(ce_loss)
correct_prediction = tf.equal(tf.argmax(self.target, 1), tf.argmax(self.pred, 1))
self.correct_labels = tf.reduce_sum(tf.cast(correct_prediction, "float"))
delta = tf.subtract(self.pred, self.target)
elif self.output_layer == 'MSE':
se_loss = 0.5 * tf.squared_difference(y, self.target)
self.loss = tf.reduce_mean(se_loss)
self.total_loss = tf.reduce_sum(se_loss)
delta = y - self.target
self.correct_labels = tf.constant(0)
else:
assert False
trainable_vars.extend([W, b_out])
h_decay = tf.subtract(1.0, tf.scalar_mul(learning_rate, tf.pow(phi, 2)))
h_decay = tf.reshape(tf.tile(h_decay, [1, tf.shape(self.h)[1]]), [-1, tf.shape(self.h)[1], tf.shape(self.h)[0]])
h_decay = tf.transpose(h_decay, [0, 2, 1])
self.h_decay = tf.reduce_sum(h_decay, axis=0)
h_delta = tf.reshape(tf.tile(delta, [1, tf.shape(self.h)[0]]), [-1, tf.shape(self.h)[0], tf.shape(self.h)[1]])
self.h_delta = tf.reduce_sum(h_delta, axis=0)
new_grads = []
phi_phi_grad = tf.multiply(phi, gate.gate_fun_gradient(phi, net))
weight = tf.transpose(tf.matmul(self.h, tf.transpose(delta)))
phi_phi_grad = tf.multiply(phi_phi_grad, weight)
new_u_grad = tf.matmul(tf.transpose(var_in), phi_phi_grad)
new_u_grad = tf.scalar_mul(1.0 / tf.cast(tf.shape(var_in)[0], tf.float32), new_u_grad)
phi_error = tf.matmul(delta, tf.transpose(W))
net_error = tf.multiply(phi_error, gate.gate_fun_gradient(phi, net))
bp_u_grad = tf.matmul(tf.transpose(var_in), net_error)
bp_u_grad = tf.scalar_mul(1.0 / tf.cast(tf.shape(var_in)[0], tf.float32), bp_u_grad)
new_u_grad = (1 - self.lam) * new_u_grad + self.lam * bp_u_grad
new_grads.append(new_u_grad)
new_b_hidden_grad = tf.reduce_mean(phi_phi_grad, axis=0)
bp_b_hidden_grad = tf.reduce_mean(net_error, axis=0)
new_b_hidden_grad = (1 - self.lam) * new_b_hidden_grad + self.lam * bp_b_hidden_grad
new_grads.append(new_b_hidden_grad)
old_grads = optimizer.compute_gradients(self.loss, var_list=[U, b_hidden])
for i, (grad, var) in enumerate(old_grads):
old_grads[i] = (new_grads[i], var)
other_grads = optimizer.compute_gradients(self.loss, var_list=trainable_vars)
self.all_gradients = old_grads + other_grads
self.train_op = optimizer.apply_gradients(self.all_gradients)
self.h_var = np.zeros((dim_hidden, dim_out))
self.variables = [W, b_out, U, b_hidden]
self.y = y
def get_assign_ops(self, src_network):
assign_ops = []
for dst_var, src_var in zip(self.variables, src_network.variables):
assign_ops.append(dst_var.assign(src_var))
return assign_ops
def predict(self, sess, x):
y = sess.run(self.y, feed_dict={self.x: x})
return y
def learn(self, sess, train_x, train_y):
_, h_decay_var, h_delta_var = \
sess.run([self.train_op, self.h_decay, self.h_delta],
feed_dict={
self.x: train_x,
self.target: train_y,
self.h: self.h_var
})
batch_size = float(train_x.shape[0])
self.h_var = np.multiply(h_decay_var / batch_size, self.h_var) - self.learning_rate * h_delta_var / batch_size
+9 -12
View File
@@ -4,18 +4,15 @@
# declaration at the top #
#######################################################################
import tensorflow as tf
from network import *
from replay import *
from policy import *
import numpy as np
class DQNAgent:
def __init__(self, name, task, network_fn, policy_fn, replay_fn, discount, step_limit, target_network_update_freq):
with tf.variable_scope(name):
self.learning_network = network_fn('learning')
self.target_network = network_fn('target')
self.assign_ops = self.target_network.get_assign_ops(self.learning_network)
def __init__(self, task, network_fn, policy_fn, replay_fn, discount, step_limit, target_network_update_freq):
self.learning_network = network_fn()
self.target_network = network_fn()
self.task = task
self.step_limit = step_limit
self.replay = replay_fn()
@@ -24,12 +21,12 @@ class DQNAgent:
self.policy = policy_fn()
self.total_steps = 0
def episode(self, sess):
def episode(self):
state = self.task.reset()
total_reward = 0.0
steps = 0
while not self.step_limit or steps < self.step_limit:
value = self.learning_network.predict(sess, np.reshape(state, (1, -1)) )
value = self.learning_network.predict(np.reshape(state, (1, -1)) )
action = self.policy.sample(value.flatten())
next_state, reward, done, info = self.task.step(action)
total_reward += reward
@@ -42,14 +39,14 @@ class DQNAgent:
experiences = self.replay.sample()
if experiences is not None:
states, actions, rewards, next_states, terminals = experiences
targets = self.learning_network.predict(sess, states)
q_next = self.target_network.predict(sess, next_states)
targets = self.learning_network.predict(states)
q_next = self.target_network.predict(next_states)
q_next = np.max(q_next, axis=1)
q_next = rewards + self.discount * q_next
q_next = np.where(terminals, 0, q_next)
targets[np.arange(len(actions)), actions] = q_next
self.learning_network.learn(sess, states, targets)
self.learning_network.learn(states, targets)
if self.total_steps % self.target_network_update_freq == 0:
sess.run(self.assign_ops)
self.target_network.sync_with(self.learning_network)
self.policy.update_epsilon()
return total_reward
+58 -37
View File
@@ -4,46 +4,67 @@
# declaration at the top #
#######################################################################
import tensorflow as tf
from common import *
import torch
from torch.autograd import Variable
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from SMDWrapper import SMDWrapper
class Network:
def __init__(self, name, dim_in, dim_out, optimizer_fn, initializer=tf.random_normal_initializer()):
self.x = tf.placeholder(tf.float32, shape=(None, dim_in))
dim_hidden1 = 50
dim_hidden2 = 200
W1, b1, net1, phi1 = \
fully_connected(name, 'layer1', self.x, dim_in, dim_hidden1, initializer, tf.nn.relu)
W2, b2, net2, phi2 = \
fully_connected(name, 'layer2', phi1, dim_hidden1, dim_hidden2, initializer, tf.nn.relu)
W3, b3, net3, self.y = \
fully_connected(name, 'layer3', phi2, dim_hidden2, dim_out, initializer, tf.identity)
self.target = tf.placeholder(tf.float32, shape=(None, dim_out))
loss = 0.5 * tf.reduce_mean(tf.squared_difference(self.y, self.target))
self.variables = [W1, b1, W2, b2, W3, b3]
self.train_op = optimizer_fn(name).minimize(loss=loss)
class FullyConnectedNet(nn.Module):
def __init__(self, dims, learning_rate, gpu=True):
super(FullyConnectedNet, self).__init__()
self.fc1 = nn.Linear(dims[0], dims[1])
self.fc2 = nn.Linear(dims[1], dims[2])
self.fc3 = nn.Linear(dims[2], dims[3])
self.criterion = nn.MSELoss()
self.learning_rate = learning_rate
self.optimizer = torch.optim.SGD(self.parameters(), learning_rate)
self.gpu = gpu and torch.cuda.is_available()
if self.gpu:
print 'Transferring network to GPU...'
self.cuda()
print 'Network transferred.'
def get_assign_ops(self, src_network):
assign_ops = []
for dst_var, src_var in zip(self.variables, src_network.variables):
assign_ops.append(dst_var.assign(src_var))
return assign_ops
def forward(self, x):
x = torch.from_numpy(np.asarray(x, dtype='float32'))
if self.gpu:
x = x.cuda()
x = Variable(x)
def predict(self, sess, x):
y = sess.run(self.y, feed_dict={self.x: x})
y = F.relu(self.fc1(x))
y = F.relu(self.fc2(y))
y = self.fc3(y)
return y
def learn(self, sess, x, target):
sess.run(self.train_op, feed_dict={self.x: x, self.target: target})
def sync_with(self, src_net):
for param_dst, param_src in zip(self.parameters(), src_net.parameters()):
param_dst.data.copy_(param_src.data)
class SimpleNetwork(Network):
def __init__(self, name, dim_in, dim_out, dim_hidden, optimizer_fn, initializer=tf.random_normal_initializer()):
self.x = tf.placeholder(tf.float32, shape=(None, dim_in))
W1, b1, net1, phi1 = \
fully_connected(name, 'layer1', self.x, dim_in, dim_hidden, initializer, tf.nn.relu)
W2, b2, net2, self.y = \
fully_connected(name, 'layer2', phi1, dim_hidden, dim_out, initializer, tf.identity)
self.target = tf.placeholder(tf.float32, shape=(None, dim_out))
loss = 0.5 * tf.reduce_mean(tf.squared_difference(self.y, self.target))
self.variables = [W1, b1, W2, b2]
self.train_op = optimizer_fn(name).minimize(loss=loss)
def predict(self, x):
return self.forward(x).cpu().data.numpy()
def learn(self, x, target):
target = torch.from_numpy(target)
if self.gpu:
target = target.cuda()
target = Variable(target)
y = self.forward(x)
loss = self.criterion(y, target)
self.zero_grad()
loss.backward()
self.optimizer.step()
def output_transfer(self, y):
return y
class SMDNetworkWrapper(SMDWrapper):
def __init__(self, net):
SMDWrapper.__init__(self, net)
def sync_with(self, src_net):
self.net.sync_with(src_net.net)
def parameters(self):
return self.net.parameters()
+20 -39
View File
@@ -33,9 +33,7 @@ class MountainCar(BasicTask):
self.env = gym.make(self.name)
self.env._max_episode_steps = sys.maxsize
optimizer_fn = lambda name: tf.train.GradientDescentOptimizer(name=name, learning_rate=0.01)
self.network_fn = lambda name: Network(name, self.state_space_size,
self.action_space_size, optimizer_fn, tf.random_normal_initializer())
self.network_fn = lambda learning_rate=0.01: FullyConnectedNet([self.state_space_size, 50, 200, self.action_space_size], learning_rate)
self.policy_fn = lambda: GreedyPolicy(epsilon=0.5, decay_factor=0.95, min_epsilon=0.1)
self.replay_fn = lambda: Replay(memory_size=10000, batch_size=10)
@@ -51,45 +49,28 @@ class CartPole(BasicTask):
def __init__(self):
self.env = gym.make(self.name)
optimizer_fn = lambda name: tf.train.GradientDescentOptimizer(name=name, learning_rate=0.01)
self.network_fn = lambda name: Network(name, self.state_space_size,
self.action_space_size, optimizer_fn, tf.random_normal_initializer())
self.network_fn = lambda learning_rate=0.01: FullyConnectedNet([self.state_space_size, 50, 200, self.action_space_size], learning_rate)
self.policy_fn = lambda: GreedyPolicy(epsilon=0.5, decay_factor=0.95, min_epsilon=0.1)
self.replay_fn = lambda: Replay(memory_size=10000, batch_size=10)
class LunarLander(BasicTask):
state_space_size = 8
action_space_size = 4
name = 'LunarLander-v2'
success_threshold = 200
discount = 0.99
step_limit = 5000
target_network_update_freq = 200
def __init__(self):
self.env = gym.make(self.name)
optimizer_fn = lambda name: tf.train.GradientDescentOptimizer(name=name, learning_rate=0.001)
self.network_fn = lambda name: Network(name, self.state_space_size,
self.action_space_size, optimizer_fn, tf.random_normal_initializer())
self.policy_fn = lambda: GreedyPolicy(epsilon=0.5, decay_factor=0.99, min_epsilon=0.1)
self.replay_fn = lambda: Replay(memory_size=20000, batch_size=100)
if __name__ == '__main__':
# task = MountainCar()
task = LunarLander()
agent = DQNAgent(task.name, task, task.network_fn, task.policy_fn, task.replay_fn,
task = MountainCar()
bp_network_fn = lambda learning_rate=0.001: FullyConnectedNet([task.state_space_size, 50, 200, task.action_space_size], learning_rate, gpu=False)
def smd_network_fn(learning_rate=0.001):
bp_network = bp_network_fn(learning_rate)
return SMDNetworkWrapper(bp_network)
agent = DQNAgent(task, smd_network_fn, task.policy_fn, task.replay_fn,
task.discount, task.step_limit, task.target_network_update_freq)
window_size = 100
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
ep = 0
rewards = []
while True:
ep += 1
reward = agent.episode(sess)
rewards.append(reward)
if len(rewards) > window_size:
reward = np.mean(rewards[-window_size:])
print 'episode %d: %f' % (ep, reward)
if reward > task.success_threshold:
break
ep = 0
rewards = []
while True:
ep += 1
reward = agent.episode()
rewards.append(reward)
if len(rewards) > window_size:
reward = np.mean(rewards[-window_size:])
print 'episode %d: %f' % (ep, reward)
if reward > task.success_threshold:
break