<p dir="auto">When dealing with discrete stochastic units, one often wants to define custom gradients. It seems like the most legitimate way to it in Tensorflow is to use <code class="notranslate">grad_func</code> keyword argument of <code class="notranslate">function.Defun</code> from <code class="notranslate">tensorflow.python.framework</code>.  However, the gradient function is only given the inputs of the function defined and the gradients, whereas it would be often very convenient to also use the outputs.</p>
<p dir="auto">For example, see the code below.</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="# In current Tensorflow I have to do like this:

@function.Defun(*(3 * [tf.float32]))
def oneovertwo_unit_bprop(probs, noise, grads):
  outcome = tf.to_float(tf.less(noise, probs))
  outcome_probs = outcome * probs + (1 - outcome) * (1 - probs)
  return grads / 2. / outcome_probs, tf.zeros_like(noise)


@function.Defun(*(2 * [tf.float32]), grad_func=oneovertwo_unit_bprop)
def oneovertwo_unit_fprop(probs, noise):
  return tf.to_float(tf.less(noise, probs))


def oneovertwo_unit(probs):
  return oneovertwo_unit_fprop(probs, tf.random_uniform(tf.shape(probs)))

# I would like to do like this

def oneovertwo_unit_bprop(probs, outcome, grads):
  outcome_probs = outcome * probs + (1 - outcome) * (1 - probs)
  return grads / 2. / outcome_probs


@function.Defun(*(2 * [tf.float32]), grad_func=oneovertwo_unit_bprop)
def oneovertwo_unit(probs):
  return tf.to_float(tf.less(tf.random_uniform(probs.get_shape()), probs))"><pre class="notranslate"><span class="pl-c"># In current Tensorflow I have to do like this:</span>

<span class="pl-en">@<span class="pl-s1">function</span>.<span class="pl-v">Defun</span>(<span class="pl-c1">*</span>(<span class="pl-c1">3</span> <span class="pl-c1">*</span> [<span class="pl-s1">tf</span>.<span class="pl-s1">float32</span>]))</span>
<span class="pl-k">def</span> <span class="pl-en">oneovertwo_unit_bprop</span>(<span class="pl-s1">probs</span>, <span class="pl-s1">noise</span>, <span class="pl-s1">grads</span>):
  <span class="pl-s1">outcome</span> <span class="pl-c1">=</span> <span class="pl-s1">tf</span>.<span class="pl-en">to_float</span>(<span class="pl-s1">tf</span>.<span class="pl-en">less</span>(<span class="pl-s1">noise</span>, <span class="pl-s1">probs</span>))
  <span class="pl-s1">outcome_probs</span> <span class="pl-c1">=</span> <span class="pl-s1">outcome</span> <span class="pl-c1">*</span> <span class="pl-s1">probs</span> <span class="pl-c1">+</span> (<span class="pl-c1">1</span> <span class="pl-c1">-</span> <span class="pl-s1">outcome</span>) <span class="pl-c1">*</span> (<span class="pl-c1">1</span> <span class="pl-c1">-</span> <span class="pl-s1">probs</span>)
  <span class="pl-k">return</span> <span class="pl-s1">grads</span> <span class="pl-c1">/</span> <span class="pl-c1">2.</span> <span class="pl-c1">/</span> <span class="pl-s1">outcome_probs</span>, <span class="pl-s1">tf</span>.<span class="pl-en">zeros_like</span>(<span class="pl-s1">noise</span>)


<span class="pl-en">@<span class="pl-s1">function</span>.<span class="pl-v">Defun</span>(<span class="pl-c1">*</span>(<span class="pl-c1">2</span> <span class="pl-c1">*</span> [<span class="pl-s1">tf</span>.<span class="pl-s1">float32</span>]), <span class="pl-s1">grad_func</span><span class="pl-c1">=</span><span class="pl-s1">oneovertwo_unit_bprop</span>)</span>
<span class="pl-k">def</span> <span class="pl-en">oneovertwo_unit_fprop</span>(<span class="pl-s1">probs</span>, <span class="pl-s1">noise</span>):
  <span class="pl-k">return</span> <span class="pl-s1">tf</span>.<span class="pl-en">to_float</span>(<span class="pl-s1">tf</span>.<span class="pl-en">less</span>(<span class="pl-s1">noise</span>, <span class="pl-s1">probs</span>))


<span class="pl-k">def</span> <span class="pl-en">oneovertwo_unit</span>(<span class="pl-s1">probs</span>):
  <span class="pl-k">return</span> <span class="pl-en">oneovertwo_unit_fprop</span>(<span class="pl-s1">probs</span>, <span class="pl-s1">tf</span>.<span class="pl-en">random_uniform</span>(<span class="pl-s1">tf</span>.<span class="pl-en">shape</span>(<span class="pl-s1">probs</span>)))

<span class="pl-c"># I would like to do like this</span>

<span class="pl-k">def</span> <span class="pl-en">oneovertwo_unit_bprop</span>(<span class="pl-s1">probs</span>, <span class="pl-s1">outcome</span>, <span class="pl-s1">grads</span>):
  <span class="pl-s1">outcome_probs</span> <span class="pl-c1">=</span> <span class="pl-s1">outcome</span> <span class="pl-c1">*</span> <span class="pl-s1">probs</span> <span class="pl-c1">+</span> (<span class="pl-c1">1</span> <span class="pl-c1">-</span> <span class="pl-s1">outcome</span>) <span class="pl-c1">*</span> (<span class="pl-c1">1</span> <span class="pl-c1">-</span> <span class="pl-s1">probs</span>)
  <span class="pl-k">return</span> <span class="pl-s1">grads</span> <span class="pl-c1">/</span> <span class="pl-c1">2.</span> <span class="pl-c1">/</span> <span class="pl-s1">outcome_probs</span>


<span class="pl-en">@<span class="pl-s1">function</span>.<span class="pl-v">Defun</span>(<span class="pl-c1">*</span>(<span class="pl-c1">2</span> <span class="pl-c1">*</span> [<span class="pl-s1">tf</span>.<span class="pl-s1">float32</span>]), <span class="pl-s1">grad_func</span><span class="pl-c1">=</span><span class="pl-s1">oneovertwo_unit_bprop</span>)</span>
<span class="pl-k">def</span> <span class="pl-en">oneovertwo_unit</span>(<span class="pl-s1">probs</span>):
  <span class="pl-k">return</span> <span class="pl-s1">tf</span>.<span class="pl-en">to_float</span>(<span class="pl-s1">tf</span>.<span class="pl-en">less</span>(<span class="pl-s1">tf</span>.<span class="pl-en">random_uniform</span>(<span class="pl-s1">probs</span>.<span class="pl-en">get_shape</span>()), <span class="pl-s1">probs</span>))</pre></div>
<p dir="auto">FYI, this an implementation of the gradient estimator from the DARN paper (<a href="https://arxiv.org/pdf/1310.8499v2.pdf" rel="nofollow">https://arxiv.org/pdf/1310.8499v2.pdf</a>).</p>
<h3 dir="auto">System information</h3>
<ul dir="auto">
<li><strong>Have I written custom code (as opposed to using a stock example script provided in TensorFlow)</strong>: Yes</li>
<li><strong>OS Platform and Distribution (e.g., Linux Ubuntu 16.04)</strong>: Linux 16.04</li>
<li><strong>TensorFlow installed from (source or binary)</strong>: source</li>
<li><strong>TensorFlow version (use command below)</strong>: 1.8.0</li>
<li><strong>Python version</strong>: 3.6</li>
<li><strong>Bazel version (if compiling from source)</strong>: 0.11</li>
<li><strong>GCC/Compiler version (if compiling from source)</strong>: gcc 5.4.0</li>
<li><strong>CUDA/cuDNN version</strong>: 9.0 - 7.1</li>
<li><strong>GPU model and memory</strong>: Titan X</li>
<li><strong>Exact command to reproduce</strong>: N/A</li>
</ul>
<p dir="auto">We are trying to use distribute.MirroredStrategy to train a model on multiple GPUs on a single machine. We have a working implementation without warm start (training from scratch) and now want to initialise to model using checkpoint from imagenet before the training.<br>
Our first attempt is to add <code class="notranslate">init_from_checkpoint</code> in our model function:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="def model_fn(features, labels, mode, params):

    .....

   tf.train.init_from_checkpoint(params['resnet_checkpoint'], {'/': 'resnet50/'})

   ...."><pre class="notranslate"><code class="notranslate">def model_fn(features, labels, mode, params):

    .....

   tf.train.init_from_checkpoint(params['resnet_checkpoint'], {'/': 'resnet50/'})

   ....
</code></pre></div>
<p dir="auto">But, this gives us the following error</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content=".../tensorflow/contrib/distribute/python/values.py&quot;, line 285, in _get_update_device
    &quot;Use DistributionStrategy.update() to modify a MirroredVariable.&quot;)"><pre class="notranslate"><code class="notranslate">.../tensorflow/contrib/distribute/python/values.py", line 285, in _get_update_device
    "Use DistributionStrategy.update() to modify a MirroredVariable.")
</code></pre></div>
<p dir="auto">Is there an example of how to warm start a training with an existing checkpoint ?</p>