def make_optimizer(env):
wrapper = optimizer.OptimizerWrapper()
lr = optimizer.make_optimizer_variable(
'learning_rate',
get_env('trainer.policy_learning_rate'),
prefix='policy_')
wrapper.set_base_optimizer(optimizer.base.AdamOptimizer(lr, epsilon=1e-3))
wrapper.append_grad_modifier(
optimizer.grad_modifier.LearningRateMultiplier([
('*/b', 2.0),
]))
env.set_policy_optimizer(wrapper)
use_linear_vr = get_env('ppo.use_linear_vr')
if not use_linear_vr:
wrapper = optimizer.OptimizerWrapper()
lr = optimizer.make_optimizer_variable(
'learning_rate',
get_env('trainer.value_learning_rate'),
prefix='value_')
wrapper.set_base_optimizer(
optimizer.base.AdamOptimizer(lr, epsilon=1e-3))
wrapper.append_grad_modifier(
optimizer.grad_modifier.LearningRateMultiplier([
('*/b', 2.0),
]))
env.set_value_optimizer(wrapper)
def make_optimizer(env):
lr = optimizer.base.make_optimizer_variable('learning_rate', get_env('trainer.learning_rate'))
wrapper = optimizer.OptimizerWrapper()
# generator learns 5 times faster
wrapper.set_base_optimizer(optimizer.base.AdamOptimizer(lr * 5., beta1=0.5, epsilon=1e-6))
wrapper.append_grad_modifier(optimizer.grad_modifier.LearningRateMultiplier([
('*/b', 2.0),
]))
env.set_g_optimizer(wrapper)
wrapper = optimizer.OptimizerWrapper()
wrapper.set_base_optimizer(optimizer.base.AdamOptimizer(lr, beta1=0.5, epsilon=1e-6))
wrapper.append_grad_modifier(optimizer.grad_modifier.LearningRateMultiplier([
('*/b', 2.0),
]))
env.set_d_optimizer(wrapper)
def make_optimizer(env):
wrapper = optimizer.OptimizerWrapper()
wrapper.set_base_optimizer(optimizer.base.MomentumOptimizer(get_env('trainer.learning_rate'), 0.9))
wrapper.append_grad_modifier(optimizer.grad_modifier.LearningRateMultiplier([
('*/b', 2.0),
]))
env.set_optimizer(wrapper)
def make_rpredictor_optimizer(env):
wrapper = optimizer.OptimizerWrapper()
wrapper.set_base_optimizer(
optimizer.base.AdamOptimizer(get_env('rpredictor.learning_rate'),
epsilon=1e-3))
wrapper.append_grad_modifier(
optimizer.grad_modifier.LearningRateMultiplier([
('*/b', 2.0),
]))
env.set_optimizer(wrapper)
def make_optimizer(env):
lr = optimizer.base.make_optimizer_variable('learning_rate', get_env('trainer.learning_rate'))
wrapper = optimizer.OptimizerWrapper()
wrapper.set_base_optimizer(optimizer.base.AdamOptimizer(lr))
wrapper.append_grad_modifier(optimizer.grad_modifier.LearningRateMultiplier([
('*/b', 2.0),
]))
env.set_g_optimizer(wrapper)
env.set_d_optimizer(wrapper)
def make_optimizer(env):
wrapper = optimizer.OptimizerWrapper()
wrapper.set_base_optimizer(optimizer.base.AdamOptimizer(get_env('trainer.learning_rate'), beta1=0.75, beta2=0.5))
wrapper.append_grad_modifier(optimizer.grad_modifier.LearningRateMultiplier([
('*/b', 2.0),
]))
# wrapper.append_grad_modifier(optimizer.grad_modifier.WeightDecay([
# ('*/W', 0.0005)
# ]))
env.set_optimizer(wrapper)
def make_optimizer(env):
lr = optimizer.base.make_optimizer_variable('learning_rate', get_env('trainer.learning_rate'))
wrapper = optimizer.OptimizerWrapper()
wrapper.set_base_optimizer(optimizer.base.AdamOptimizer(lr, epsilon=1e-3))
wrapper.append_grad_modifier(optimizer.grad_modifier.LearningRateMultiplier([
('*/b', 2.0),
]))
wrapper.append_grad_modifier(optimizer.grad_modifier.GlobalGradClipByAvgNorm(0.1))
env.set_optimizer(wrapper)
def make_optimizer(env):
lr = optimizer.base.make_optimizer_variable(
'learning_rate', get_env('trainer.learning_rate'))
wrapper = optimizer.OptimizerWrapper()
wrapper.set_base_optimizer(
optimizer.base.RMSPropOptimizer(lr, epsilon=1e-3))
wrapper.append_grad_modifier(
optimizer.grad_modifier.LearningRateMultiplier([
('*/b', 2.0),
]))
wrapper.append_grad_modifier(
optimizer.grad_modifier.WeightDecay([
('*/W', 5e-4),
]))
env.set_optimizer(wrapper)
def make_optimizer(env):
opt = rl.train.TRPOOptimizer(env,
max_kl=get_env('trpo.max_kl'),
cg_damping=get_env('trpo.cg.damping'))
env.set_policy_optimizer(opt)
use_linear_vr = get_env('trpo.use_linear_vr')
if not use_linear_vr:
wrapper = optimizer.OptimizerWrapper()
wrapper.set_base_optimizer(
optimizer.base.AdamOptimizer(
get_env('trainer.value_learning_rate'), epsilon=1e-3))
wrapper.append_grad_modifier(
optimizer.grad_modifier.LearningRateMultiplier([
('*/b', 2.0),
]))
env.set_value_optimizer(wrapper)
def main():
from tartist.plugins.trainer_enhancer import snapshot
img = image.imread(args.image_path)
smg = image.imread(args.style_path)
h, w = img.shape[0:2]
env = Env(master_dev=args.device)
with env.as_default():
make_network(env)
func = env.make_func()
func.compile(env.network.outputs)
env.initialize_all_variables()
snapshot.load_weights_file(env, args.weight_path)
res_img = func(img=img[np.newaxis])
res_smg = func(img=smg[np.newaxis])
# create a new env for train
env = SimpleTrainerEnv(master_dev=args.device)
with env.as_default():
make_network(env, h, w)
net = env.network
netin = net.outputs['img']
with env.as_default():
outputs = net.outputs
loss_content = 0.
for i in get_env('neural_style.content_layers'):
loss_content += i[1] * neural_style_opr.get_content_loss(res_img[i[0]], outputs[i[0]])
loss_style = 0.
for i in get_env('neural_style.style_layers'):
loss_style += i[1] * neural_style_opr.get_style_loss(res_smg[i[0]], outputs[i[0]])
loss = loss_content + loss_style * get_env('neural_style.style_strenth')
net.set_loss(loss)
wrapper = optimizer.OptimizerWrapper()
wrapper.set_base_optimizer(optimizer.base.AdamOptimizer(get_env('neural_style.learning_rate')))
wrapper.append_grad_modifier(optimizer.grad_modifier.LearningRateMultiplier([
('img', 1.0),
('*', 0.0),
]))
env.set_optimizer(wrapper)
func = env.make_optimizable_func()
func.compile({'loss': net.loss})
env.initialize_all_variables()
snapshot.load_weights_file(env, args.weight_path)
noise_img = np.random.uniform(-20, 20, (h, w, 3)).astype('float32')
image_mean = get_env('neural_style.image_mean')
image_noise_ratio = get_env('neural_style.image_noise_ratio')
jmg = (1 - image_noise_ratio) * (img - image_mean) + image_noise_ratio * noise_img + image_mean
netin.set_value(jmg[np.newaxis])
io.makedir(args.output_path)
for i in range(args.nr_iters + 1):
if i != 0:
r = func()
logger.info('iter {}: loss = {}'.format(i, float(r['loss'])))
if i % args.save_step == 0:
output = netin.get_value()[0]
output_path = os.path.join(args.output_path, 'iter_{:04d}.jpg'.format(i))
image.imwrite(output_path, output)
logger.critical('Output written: {}'.format(output_path))