本文整理汇总了Python中utils.average_gradients方法的典型用法代码示例。如果您正苦于以下问题：Python utils.average_gradients方法的具体用法？Python utils.average_gradients怎么用？Python utils.average_gradients使用的例子？那么, 这里精选的方法代码示例或许可以为您提供帮助。您也可以进一步了解该方法所在类utils的用法示例。
在下文中一共展示了utils.average_gradients方法的2个代码示例，这些例子默认根据受欢迎程度排序。您可以为喜欢或者感觉有用的代码点赞，您的评价将有助于系统推荐出更棒的Python代码示例。
示例1: step
# 需要导入模块: import utils [as 别名]
# 或者: from utils import average_gradients [as 别名]
def step(self):
cmp_output = self.model(self.image_input, self.sparse_input)
loss_flow = self.flow_criterion(cmp_output, self.flow_target) / self.world_size
self.optim.zero_grad()
loss_flow.backward()
utils.average_gradients(self.model)
self.optim.step()
return {'loss_flow': loss_flow} 示例2: create_train_op
# 需要导入模块: import utils [as 别名]
# 或者: from utils import average_gradients [as 别名]
def create_train_op(self, optim, iterator, global_step, regularizer_scale=1e-4, train=True, trainable=True, is_scale=True, training_mode='no_distillation'):
if self.num_gpus == 1:
losses, regularizer_loss = self.build(iterator, regularizer_scale=regularizer_scale, train=train, trainable=trainable, is_scale=is_scale, training_mode=training_mode)
optim_loss = losses['abs_robust_mean']['no_occlusion']
train_op = optim.minimize(optim_loss, var_list=tf.trainable_variables(), global_step=global_step)
else:
tower_grads = []
tower_losses = []
tower_regularizer_losses = []
with tf.variable_scope(tf.get_variable_scope()):
for i in range(self.num_gpus):
with tf.device('/gpu:%d' % i):
with tf.name_scope('tower_{}'.format(i)) as scope:
losses_, regularizer_loss_ = self.build(iterator, regularizer_scale=regularizer_scale, train=train, trainable=trainable, is_scale=is_scale, training_mode=training_mode)
optim_loss = losses_['abs_robust_mean']['no_occlusion']
# Reuse variables for the next tower.
tf.get_variable_scope().reuse_variables()
grads = self.optim.compute_gradients(optim_loss, var_list=tf.trainable_variables())
tower_grads.append(grads)
tower_losses.append(losses_)
tower_regularizer_losses.append(regularizer_loss_)
#self.add_loss_summary(losses_, keys=['abs_robust_mean', 'census'], prefix='tower_%d' % i)
grads = average_gradients(tower_grads)
train_op = optim.apply_gradients(grads, global_step=global_step)
losses = tower_losses[0].copy()
for key in losses.keys():
for loss_key, loss_value in losses[key].items():
for i in range(1, self.num_gpus):
losses[key][loss_key] += tower_losses[i][key][loss_key]
losses[key][loss_key] /= self.num_gpus
regularizer_loss = 0.
for i in range(self.num_gpus):
regularizer_loss += tower_regularizer_losses[i]
regularizer_loss /= self.num_gpus
self.add_loss_summary(losses, keys=losses.keys())
tf.summary.scalar('regularizer_loss', regularizer_loss)
return train_op, losses, regularizer_loss