def compute_apply(self):
fetches = run_timeline(
self.sess,
[self.models[0].loss, self.apply_op, self.nccl_control_out],
feed_dict=self._grad_feed_dict(),
name="compute_apply")
return fetches[0]
def ps_compute_apply(self,
out_grad_shard_oids,
agg_grad_shard_oids,
tl_name="ps_compute_apply",
write_timeline=False):
feed_dict = {
ph: oid
for (ph,
oid) in zip(self.plasma_in_grads_oids, out_grad_shard_oids)
}
feed_dict.update({
ph: oid
for (ph,
oid) in zip(self.plasma_out_grads_oids, agg_grad_shard_oids)
})
fetch(agg_grad_shard_oids)
run_timeline(
self.sess,
[self.plasma_in_grads, self.apply_op, self.nccl_control_out],
feed_dict=feed_dict,
write_timeline=write_timeline)
def ps_compute_apply(self,
out_grad_shard_oids,
agg_grad_shard_oids,
tl_name="ps_compute_apply",
write_timeline=False):
feed_dict = dict(zip(self.plasma_in_grads_oids, out_grad_shard_oids))
feed_dict.update(
dict(zip(self.plasma_out_grads_oids, agg_grad_shard_oids)))
fetch(agg_grad_shard_oids)
fetches = run_timeline(self.sess, [
self.models[0].loss, self.plasma_in_grads, self.apply_op,
self.nccl_control_out
],
feed_dict=feed_dict,
write_timeline=write_timeline)
return fetches[0]