jit::script::Module parameters are not updating when training

username: markl
cooked: <p>I’m trying to train a <code>torch::jit::script::Module</code> but the loss is not decreasing. Here is a minimal example.</p>
<pre><code class="lang-auto">#include &lt;torch/script.h&gt;
#include &lt;torch/torch.h&gt;
#include &lt;iostream&gt;
#include &lt;vector&gt;

// custom loader code
#include "nets/nets.h"
#include "util/runfiles.h"

int main(int argc, char** argv) {
  std::cout &lt;&lt; "Nets example" &lt;&lt; std::endl;

  // Custom code that loads the module on CUDA
  auto runfiles = MakeRunfiles(argv[0]);
  torch::jit::script::Module script_module = LoadSegnetBackbone(*runfiles);
  script_module.train();
  std::cout &lt;&lt; "Loaded script module" &lt;&lt; std::endl;

  // Pull parameters out of the script module so we can push them into the
  // optimizer.
  std::vector&lt;at::Tensor&gt; parameters;
  for (const auto&amp; parameter : script_module.get_parameters()) {
    parameters.push_back(parameter.value().toTensor());
  }
  torch::optim::SGD optimizer(std::move(parameters), /*lr=*/0.01);

  constexpr int kBatchSize = 1;
  for (int epoch = 1; epoch &lt;= 1000; ++epoch) {
    optimizer.zero_grad();

    // The input is a (kBatchSize,3,300,300) tensor filled with ones
    at::Tensor input = torch::ones({kBatchSize, /*channels (rgb) =*/3,
                                    /*height=*/300, /*width=*/300})
                           .to(at::kFloat)
                           .to(at::kCUDA);

    // Push the input through the script module
    std::vector&lt;torch::jit::IValue&gt; inputs;
    inputs.push_back(input);
    at::Tensor script_module_forward = script_module.forward(inputs).toTensor();
    // The result is an output tensor of size (kBatchSize, 32, 300, 300)

    // ground truth is a (kBatchSize, 300, 300) tensor filled with ones
    at::Tensor ground_truth =
        torch::ones({kBatchSize, /*height=*/300, /*width=*/300})
            .to(at::kLong)
            .to(at::kCUDA);

    at::Tensor loss = torch::nll_loss2d(
        torch::log_softmax(script_module_forward, /*dim=*/1), ground_truth);
    loss.backward();
    optimizer.step();

    if (epoch % 50 == 0) {
      std::cout &lt;&lt; "Loss was " &lt;&lt; loss.item&lt;float&gt;() &lt;&lt; std::endl;
    }
  }
}
</code></pre>
<p>And the output is</p>
<pre><code class="lang-auto">Nets example
Loaded script module
Loss was 3.44751
Loss was 3.44751
Loss was 3.44751
Loss was 3.44751
Loss was 3.44751
Loss was 3.44751
Loss was 3.44751
Loss was 3.44751
...
</code></pre>
<p>Am I doing something wrong here?</p>

username: albanD
cooked: <p>I’m not sure about the c++ semantic but are you sure that <code>parameters.push_back(parameter.value().toTensor());</code> is actually storing a reference to the parameter or is it storing a clone of it?</p>

username: markl
cooked: <p>Oh that’s a good point. Do you know if there is any resource that explains the correct way to put the parameters into the optimizer?</p>

username: tom
cooked: <pre><code class="lang-auto">  torch::optim::SGD optimizer(
      model.parameters(), torch::optim::SGDOptions(0.01).momentum(0.5));
</code></pre>
<p>(from Peter Goldsborough’s example repo).</p>

username: markl
cooked: <p>Thanks <a class="mention" href="/u/tom">@tom</a>, but I don’t think the <code>torch::jit::script::Module</code> has a <code>.parameters()</code> member function.</p>

username: tom
cooked: <p>Right, sorry. Does at::Tensor --&gt; torch::Tensor or even autograd::Variable help?</p>

username: markl
cooked: <p>Thanks for the suggestion. <code>at::Tensor</code>-&gt;<code>torch::Tensor</code> didn’t help. I’m not sure where to use <code>autograd::Variable</code>.</p>

username: tom
cooked: <p>In the same place. (Just a guess from looking at <a href="https://github.com/pytorch/pytorch/blob/2cc0f1bbc685e6d8bb1ad74a119282bf732d8f28/torch/csrc/jit/script/module.cpp#L127" rel="nofollow noopener">jit::script::Module</a>).</p>

username: markl
cooked: <p>I got a compile error from doing that since the optimizer’s constructor takes a tensor</p>
<pre><code class="lang-auto">bazel-out/k8-fastbuild/bin/third_party/libtorch/_virtual_includes/libtorch_csrc/torch/optim/sgd.h: In instantiation of 'torch::optim::SGD::SGD(ParameterContainer&amp;&amp;, const torch::optim::SGDOptions&amp;) [with ParameterContainer = std::vector&lt;torch::autograd::Variable&gt;]':
nets/nets_example_2.cpp:25:65:   required from here
bazel-out/k8-fastbuild/bin/third_party/libtorch/_virtual_includes/libtorch_csrc/torch/optim/sgd.h:36:24: error: no matching function for call to 'torch::optim::Optimizer::Optimizer(std::vector&lt;torch::autograd::Variable&gt;)'
         options(options) {}
                        ^
In file included from bazel-out/k8-fastbuild/bin/third_party/libtorch/_virtual_includes/libtorch_csrc/torch/optim/adagrad.h:4,
                 from bazel-out/k8-fastbuild/bin/third_party/libtorch/_virtual_includes/libtorch_csrc/torch/optim.h:3,
                 from bazel-out/k8-fastbuild/bin/third_party/libtorch/_virtual_includes/libtorch_csrc/torch/all.h:7,
                 from bazel-out/k8-fastbuild/bin/third_party/libtorch/_virtual_includes/libtorch_csrc/torch/torch.h:3,
                 from nets/nets_example_2.cpp:2:
bazel-out/k8-fastbuild/bin/third_party/libtorch/_virtual_includes/libtorch_csrc/torch/optim/optimizer.h:37:12: note: candidate: 'torch::optim::detail::OptimizerBase::OptimizerBase(std::vector&lt;at::Tensor&gt;)'
   explicit OptimizerBase(std::vector&lt;Tensor&gt; parameters);

</code></pre>

username: abtion
cooked: <p>hi, did you solved it? i have the same problem.</p>

username: markl
cooked: <p>I have not solved it yet. There is an open issue for it on github <a href="https://github.com/pytorch/pytorch/issues/28478" rel="nofollow noopener">https://github.com/pytorch/pytorch/issues/28478</a> maybe you can give a thumbs up on it to give it more attention.</p>
<p>I had an idea which was manually call the set parameters function after the gradient step. I haven’t tried it yet.</p>

username: abtion
cooked: <p>Thanks, I have not tried it yet. There are some codes of load script model into  torch::nn::Module on github <a href="https://github.com/kerry-Cho/transfer-learning-Libtorch" rel="nofollow noopener">https://github.com/kerry-Cho/transfer-learning-Libtorch</a>, maybe help you in some ways.</p>

username: tancl
cooked: <p>Interesting. If I understand correctly, one way of training the jit model is by defining the identical header file in C and then load the jit model with load instead of jit load.</p>
<p>However, I am also suffering by just using the jit model and train it. The example that <a class="mention" href="/u/markl">@markl</a> shows loaded the parameters, from jit model to the optimizer, and then update it with optimizer. I am wondering whether this could be done as well and whether this is the direction of pytorch.</p>
<p>If this could be done, I think it will be great as any other software could just call the libtorch and perform training and inferencing, making it outstanding compare to tensorflow.</p>
<p>just my 2 cents, apologize me if I have made any statements wrong above.</p>
<p>rgds,<br>
CL</p>

username: wayne391
cooked: <p>Same problem here.</p>
<p>I also found a similar discussion:</p><aside class="quote" data-post="1" data-topic="35932">
  <div class="title">
    <div class="quote-controls"></div>
    <img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/a/ea666f/48.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/can-scriptmodule-be-used-for-training/35932">Can ScriptModule be used for training?</a> 
  </div>
  <blockquote>
    Can ScriptModule be used for training? The reason we ask this question is that we found some source code related back propagation in torch.jit. <a href="https://github.com/pytorch/pytorch/blob/master/torch/csrc/jit/autodiff.cpp#L1151" rel="noopener nofollow ugc">https://github.com/pytorch/pytorch/blob/master/torch/csrc/jit/autodiff.cpp#L1151</a>
  </blockquote>
</aside>

<p>Any update?</p>