Getting leaf nodes of graph directly from torch.autograd.graph.Node

username: lk1422
cooked: <p>Hello I am trying to directly access specific tensors which are located in the computation graph, however each grad_fn/Node seems to store temporary tensors differently:</p>
<blockquote>
<blockquote>
<blockquote>
<p>b = lin(a)<br>
dir(b.grad_fn)<br>
[‘<strong>call</strong>’, ‘<strong>class</strong>’, ‘<strong>delattr</strong>’, ‘<strong>dir</strong>’, ‘<strong>doc</strong>’, ‘<strong>eq</strong>’, ‘<strong>format</strong>’, ‘<strong>ge</strong>’, ‘<strong>getattribute</strong>’, ‘<strong>gt</strong>’, ‘<strong>hash</strong>’, ‘<strong>init</strong>’, ‘<strong>init_subclass</strong>’, ‘<strong>le</strong>’, ‘<strong>lt</strong>’, ‘<strong>ne</strong>’, ‘<strong>new</strong>’, ‘<strong>reduce</strong>’, ‘<strong>reduce_ex</strong>’, ‘<strong>repr</strong>’, ‘<strong>setattr</strong>’, ‘<strong>sizeof</strong>’, ‘<strong>str</strong>’, ‘<strong>subclasshook</strong>’, ‘_raw_saved_mat1’, ‘_raw_saved_mat2’, ‘_register_hook_dict’, ‘_saved_alpha’, ‘_saved_beta’, ‘_saved_mat1’, ‘_saved_mat1_sym_sizes’, ‘_saved_mat1_sym_strides’, ‘_saved_mat2’, ‘_saved_mat2_sym_sizes’, ‘_saved_mat2_sym_strides’, ‘metadata’, ‘name’, ‘next_functions’, ‘register_hook’, ‘register_prehook’, ‘requires_grad’]</p>
</blockquote>
</blockquote>
</blockquote>
<p>The tensor input a is stored as _saved_mat_1</p>
<p>However when for this example</p>
<blockquote>
<blockquote>
<blockquote>
<p>c = torch.randn(3,32,32)<br>
conv = torch.nn.Conv2d(3,16,1)<br>
d = conv(c)<br>
dir(d.grad_fn.next_functions[0][0])<br>
[‘<strong>call</strong>’, ‘<strong>class</strong>’, ‘<strong>delattr</strong>’, ‘<strong>dir</strong>’, ‘<strong>doc</strong>’, ‘<strong>eq</strong>’, ‘<strong>format</strong>’, ‘<strong>ge</strong>’, ‘<strong>getattribute</strong>’, ‘<strong>gt</strong>’, ‘<strong>hash</strong>’, ‘<strong>init</strong>’, ‘<strong>init_subclass</strong>’, ‘<strong>le</strong>’, ‘<strong>lt</strong>’, ‘<strong>ne</strong>’, ‘<strong>new</strong>’, ‘<strong>reduce</strong>’, ‘<strong>reduce_ex</strong>’, ‘<strong>repr</strong>’, ‘<strong>setattr</strong>’, ‘<strong>sizeof</strong>’, ‘<strong>str</strong>’, ‘<strong>subclasshook</strong>’, ‘_raw_saved_input’, ‘_raw_saved_weight’, ‘_register_hook_dict’, ‘_saved_bias_sym_sizes_opt’, ‘_saved_dilation’, ‘_saved_groups’, ‘_saved_input’, ‘_saved_output_padding’, ‘_saved_padding’, ‘_saved_stride’, ‘_saved_transposed’, ‘_saved_weight’, ‘metadata’, ‘name’, ‘next_functions’, ‘register_hook’, ‘register_prehook’, ‘requires_grad’]</p>
</blockquote>
</blockquote>
</blockquote>
<p>Now the tensor input is stored under _saved_tensor.</p>
<p>For educational purposes I am trying to develop a system that allows distributed (model parallel) training across remote devices. So I am storing disconnected nn.Modules on each of the devices, however in order to complete the backward pass I must pass the gradients between the devices, so I need a way to easily access the input’s gradient. If there is an easier way to re collect the inputs other than traversing through the computation graph that would also help, however I am trying to make it so I have to change as little as possible about already existing nn.Modules and thus do not want to make any changes to there forward passes.</p>
<p>Thanks (sorry for any typos)</p>