def initialize(self,
inputs,
loss,
target,
priority_expr,
givens=None,
lr_mult=1):
self._target = target
params = target.get_params(trainable=True)
gradients = theano.grad(loss, wrt=params, disconnected_inputs="ignore")
if self._scale_conv_grads: # (for dueling network architecture)
gradients = scale_conv_gradients(params,
gradients,
scale=2**(-1 / 2))
gradients, grad_norm = apply_grad_norm_clip(gradients,
self._grad_norm_clip)
lr = self._learning_rate * lr_mult # (lr_mult can be shared variable)
updates = self._update_method(gradients, params, learning_rate=lr)
self._f_opt = ext.compile_function(
inputs=inputs,
outputs=[priority_expr, loss],
updates=updates,
givens=givens,
log_name="grad_and_update",
)
def initialize(self, inputs, loss, target, priority_expr,
givens=None, lr_mult=1):
self._target = target
params = target.get_params(trainable=True)
gradients = theano.grad(loss, wrt=params, disconnected_inputs="ignore")
if self._scale_conv_grads:
gradients = scale_conv_gradients(params, gradients,
scale=2 ** (-1 / 2))
# Compute gradient and save to GPU vector.
flat_grad, shared_grad, flat_update = flat_shared_grad(target, gradients)
self._shared_grad = shared_grad
# All-reduce gradient in-place in shared_grad, then reshape
gradients, avg_factor_var = avg_grads_from_flat(shared_grad, params)
self._avg_factor_var = avg_factor_var
gradients, grad_norm = apply_grad_norm_clip(gradients, self._grad_norm_clip)
lr = self._learning_rate * lr_mult # (lr_mult can be shared variable)
updates = self._update_method(gradients, params, learning_rate=lr)
self._f_gradient = ext.compile_function(
inputs=inputs,
outputs=[priority_expr, loss],
updates=[flat_update],
givens=givens,
log_name="gradient",
)
self._f_update = ext.compile_function(
inputs=[],
updates=updates,
log_name="update",
)
def initialize(self, inputs, losses, constraints, target,
givens=None, lr_mult=1):
self._target = target
loss = sum(losses)
params = target.get_params(trainable=True)
gradients = theano.grad(loss, wrt=params, disconnected_inputs='ignore')
# Phase 1: Compute gradient and save to GPU vector
flat_grad, shared_grad, flat_update = flat_shared_grad(target, gradients)
self._shared_grad = shared_grad
# Phase 2: All-reduce gradient in-place in shared_grad, then reshape
gradients, avg_factor_var = avg_grads_from_flat(shared_grad, params)
self._avg_factor_var = avg_factor_var # (set later as 1 / n_gpu)
# Phase 3: Apply combined gradient locally
gradients, grad_norm = apply_grad_norm_clip(gradients, self._grad_norm_clip)
lr = self._learning_rate * lr_mult # (lr_mult can be shared variable)
updates = self._update_method(gradients, params, learning_rate=lr)
self._f_grad = ext.compile_function(
inputs=inputs,
outputs=loss,
updates=[flat_update],
givens=givens,
log_name="gradient",
)
self._f_update = ext.compile_function(
inputs=[],
outputs=grad_norm,
updates=updates,
log_name="update",
)
def initialize(self,
inputs,
losses,
constraints,
target,
givens=None,
lr_mult=1):
self._target = target
loss = sum(losses)
params = target.get_params(trainable=True)
gradients = theano.grad(loss, wrt=params, disconnected_inputs='ignore')
gradients, grad_norm = apply_grad_norm_clip(gradients,
self._grad_norm_clip)
lr = self._learning_rate * lr_mult # (lr_mult can be shared variable)
updates = self._update_method(gradients, params, learning_rate=lr)
# Prepare to load data onto GPU (and shuffle indexes there).
load_updates, givens, opt_inputs = make_shared_inputs(
inputs, self._shuffle)
self._f_load = ext.compile_function(
inputs=inputs,
updates=load_updates,
log_name="load",
)
self._f_opt = ext.compile_function(
inputs=opt_inputs,
outputs=[loss, grad_norm],
updates=updates,
givens=givens,
log_name="grad_and_update",
)
def initialize(self, inputs, losses, constraints, target, lr_mult=1):
self._target = target
loss = sum(losses)
params = target.get_params(trainable=True)
gradients = theano.grad(loss, wrt=params, disconnected_inputs='ignore')
gradients, grad_norm = apply_grad_norm_clip(gradients, self._grad_norm_clip)
# Phase 1: Compute gradient and save to GPU vector
flat_grad, shared_grad, flat_update = flat_shared_grad(target, gradients)
# Phase 2: apply gradient chunks to update central params; e.g. rmsprop
lr = self._learning_rate * lr_mult
if self.n_update_chunks > 1:
updates_args = (shared_grad, lr, self.n_update_chunks, self._update_method_args)
chunk_inputs, outputs_list, updates, idxs = \
chunked_updates(self._update_method_name, updates_args)
self._chunk_idxs = idxs
else:
whole_inputs, whole_outputs, whole_updates = \
whole_update(self._update_method_name, shared_grad, lr, self._update_method_args)
# Phase 3: copy new param values from shared_grad to params
copy_updates = copy_params_from_flat(params, shared_grad)
# Phase 1
self._f_gradient = ext.compile_function(
inputs=inputs,
outputs=[loss, grad_norm],
updates=[flat_update],
log_name="gradient",
)
# Phase 2
if self.n_update_chunks > 1:
f_update_chunks = list()
for i, (outputs, update) in enumerate(zip(outputs_list, updates)):
f_update_chunks.append(ext.compile_function(
inputs=chunk_inputs,
outputs=outputs,
updates=[update],
log_name="update_chunk_{}".format(i))
)
self._f_update_chunks = f_update_chunks
else:
self._f_update = ext.compile_function(
inputs=whole_inputs,
outputs=whole_outputs,
updates=whole_updates,
log_name="update",
)
# Phase 3
self._f_copy = ext.compile_function(
inputs=[],
updates=copy_updates,
log_name="copy_params",
)
def initialize(self,
inputs,
loss,
target,
priority_expr,
givens=None,
lr_mult=1):
self._target = target
params = target.get_params(trainable=True)
gradients = theano.grad(loss, wrt=params, disconnected_inputs="ignore")
if self._scale_conv_grads:
gradients = scale_conv_gradients(params,
gradients,
scale=2**(-1 / 2))
# if self._layerwise_stats:
# self._n_params = len(params)
# param_grad_norms = [T.sqrt(T.sum(g ** 2)) for g in gradients]
gradients, grad_norm = apply_grad_norm_clip(gradients,
self._grad_norm_clip)
lr = self._learning_rate * lr_mult # (lr_mult can be shared variable)
updates, steps = self._update_method(gradients,
params,
learning_rate=lr)
# step_norm = T.sqrt(sum(T.sum(s ** 2) for s in steps))
# if self._layerwise_stats:
# param_step_norms = [T.sqrt(T.sum(s ** 2)) for s in steps]
# outputs = [priority_expr, loss, grad_norm, step_norm]
outputs = [priority_expr, loss]
# if self._layerwise_stats:
# outputs += param_grad_norms + param_step_norms
self._f_opt = ext.compile_function(
inputs=inputs,
outputs=outputs,
updates=updates,
givens=givens,
log_name="grad_and_update",
)
def initialize(self,
inputs,
loss,
target,
priority_expr,
givens=None,
lr_mult=1):
self._target = target
params = target.get_params(trainable=True)
gradients = theano.grad(loss, wrt=params, disconnected_inputs="ignore")
if self._scale_conv_grads:
gradients = scale_conv_gradients(params,
gradients,
scale=2**(-1 / 2))
gradients, grad_norm = apply_grad_norm_clip(gradients,
self._grad_norm_clip)
# Phase 1: Compute gradient and save to GPU vector
flat_grad, shared_grad, flat_update = flat_shared_grad(
target, gradients)
# Phase 2: apply gradient chunks to central params
lr = self._learning_rate * lr_mult
updates_args = (shared_grad, lr, self._n_update_chunks,
self._update_method_args)
chunk_inputs, outputs_list, updates, idxs = \
chunked_updates(self._update_method_name, updates_args)
self._chunk_idxs = idxs
# Phase 3: copy new param values from sared_grad to params
copy_updates = copy_params_from_flat(params, shared_grad)
if self._update_method_name == "adam":
copy_updates.append(updates.pop()) # (Move the t update)
# Phase 1
self._f_gradient = ext.compile_function(
inputs=inputs,
outputs=[priority_expr, loss, grad_norm],
updates=[flat_update],
givens=givens,
log_name="gradient",
)
# Phase 2
f_update_chunks = list()
for i, (outputs, update) in enumerate(zip(outputs_list, updates)):
f_update_chunks.append(
ext.compile_function(inputs=chunk_inputs,
outputs=outputs,
updates=[update],
log_name="update_chunk_{}".format(i)))
self._f_update_chunks = f_update_chunks
# Phase 3
self._f_copy = ext.compile_function(
inputs=[],
updates=copy_updates,
log_name="copy_params",
)
def initialize(self, inputs, losses, constraints, target, batch_size, givens=None, lr_mult=1):
self._target = target
params = target.get_params(trainable=True)
loss = T.mean(losses)
gradients = theano.grad(loss, wrt=params, disconnected_inputs='ignore')
# sep_grads = [theano.grad(losses[d], wrt=params, disconnected_inputs='ignore') for d in range(batch_size)]
# lyr_sep_grads = [list() for _ in range(len(params))]
# for d_grads in sep_grads:
# for d_grad, lyr_sep_grad in zip(d_grads, lyr_sep_grads):
# lyr_sep_grad.append(d_grad)
# lyr_sep_grads = [T.stack(lsg) for lsg in lyr_sep_grads]
# lyr_sep_grd_sqnorms = [T.sum(lsg ** 2) for lsg in lyr_sep_grads]
# lyr_cmb_grd_sqnorms = [T.sum(T.sum(lsg, axis=0) ** 2) for lsg in lyr_sep_grads]
# lyr_grad_diversities = [sep / cmb for sep, cmb in
# zip(lyr_sep_grd_sqnorms, lyr_cmb_grd_sqnorms)]
# tot_sep_grd_sqnorms = T.sum(lyr_sep_grd_sqnorms)
# tot_cmb_grd_sqnorms = T.sum(lyr_cmb_grd_sqnorms)
# grad_diversity = tot_sep_grd_sqnorms / tot_cmb_grd_sqnorms
# batch_size_bound = batch_size * grad_diversity
# # Jacobians broke inside Theano.
# # jacobians = theano.gradient.jacobian(losses, wrt=params, disconnected_inputs='ignore')
# # lyr_sep_grd_sqnorms = [T.sum(j ** 2) for j in jacobians]
# # lyr_cmb_grd_sqnorms = [T.sum(T.sum(j, axis=0) ** 2) for j in jacobians]
# # lyr_grad_diversities = [sep / cmb for sep, cmb in
# # zip(lyr_sep_grd_sqnorms, lyr_cmb_grd_sqnorms)]
# # tot_sep_grd_sqnorms = T.sum(lyr_sep_grd_sqnorms)
# # tot_cmb_grd_sqnorms = T.sum(lyr_cmb_grd_sqnorms)
# # grad_diversity = tot_sep_grd_sqnorms / tot_cmb_grd_sqnorms
# # n = inputs[0].shape[0]
# # batch_size_bound = n * grad_diversity
# check_cmb_grd_sqnorms = sum(T.sum(g ** 2) for g in gradients) * batch_size
if self._layerwise_stats:
self._n_params = len(params)
param_grad_norms = [T.sqrt(T.sum(g ** 2)) for g in gradients]
gradients, grad_norm = apply_grad_norm_clip(gradients, self._grad_norm_clip)
lr = self._learning_rate * lr_mult # (lr_mult can be theano shared variable)
updates, steps = self._update_method(gradients, params, learning_rate=lr)
step_norm = T.sqrt(sum(T.sum(s ** 2) for s in steps))
if self._layerwise_stats:
param_step_norms = [T.sqrt(T.sum(s ** 2)) for s in steps]
outputs = [grad_norm, step_norm]
# grad_div_outputs = [grad_diversity, batch_size_bound]
if self._layerwise_stats:
outputs += param_grad_norms + param_step_norms
# grad_div_outputs += lyr_grad_diversities
# ipdb.set_trace()
self._opt_fun["optimize"] = ext.compile_function(
inputs=inputs,
outputs=outputs,
updates=updates,
givens=givens,
log_name="grad_and_update",
)
self._opt_fun["gradient"] = ext.compile_function(
inputs=inputs,
outputs=gradients,
log_name="gradients_only",
)
