Trying to backward through the graph a second time when using constant tensors

magicaltoast: <p>H I want to tweak aladdinpersson implementation of style transfer.<br>
I have two tensors that are not changing - features of original and style image.<br>
But in the original implementation, I compute them on every step.</p>
<p>This looks like this</p>
<pre data-code-wrap="py"><code class="lang-nohighlight">for step in trange(6_000):
    generated_features = model(generated)

    # Two following tensors are constant
    original_img_features = model(original_img) 
    style_features = model(style_img)

    style_loss = original_loss = 0

    for gen_feature, orig_feature, style_feature in zip(
        generated_features, original_img_features, style_features
    ):
        batch_size, channel, height, width = gen_feature.shape
        original_loss += torch.mean((gen_feature - orig_feature) ** 2)

        G = gen_feature.view(channel, height * width).mm(
            gen_feature.view(channel, height * width).t()
        )

        A = style_feature.view(channel, height * width).mm(
            style_feature.view(channel, height * width).t()
        )
        style_loss += torch.mean((G - A) ** 2)

    total_loss = alpha * original_loss + beta * style_loss
    optimizer.zero_grad()
    total_loss.backward(retain_graph=True)
    optimizer.step()
</code></pre>
<p>When I move the creation of these tensors outside the loop I get an <code>Trying to backward through the graph a second time (or directly access saved tensors after they have already been freed). Saved intermediate values of the graph are freed when you call .backward() or autograd.grad(). Specify retain_graph=True if you need to backward through the graph a second time or if you need to access saved tensors after calling backward.</code> error. How can I create these tensors and retain them during the whole training? I am pretty sure that I don’t need saving the whole graph, just these two tensors</p>
<pre data-code-wrap="py"><code class="lang-nohighlight">original_img_features = model(original_img) 
style_features = model(style_img)

for step in trange(6_000):
    generated_features = model(generated)
    style_loss = original_loss = 0

    for gen_feature, orig_feature, style_feature in zip(
        generated_features, original_img_features, style_features
    ):
        batch_size, channel, height, width = gen_feature.shape
        original_loss += torch.mean((gen_feature - orig_feature) ** 2)

        G = gen_feature.view(channel, height * width).mm(
            gen_feature.view(channel, height * width).t()
        )

        A = style_feature.view(channel, height * width).mm(
            style_feature.view(channel, height * width).t()
        )
        style_loss += torch.mean((G - A) ** 2)

    total_loss = alpha * original_loss + beta * style_loss
    optimizer.zero_grad()
    total_loss.backward(retain_graph=True)
    optimizer.step()
</code></pre>

Dazitu616: <p>Do you want the gradient to also BP through <code>original_img_features</code> and <code>style_features</code>? If yes, then you will have to use <code>retain_graph=True</code>. If not (i.e. the gradient is only related to generated_features), I think you can do something like this:</p>
<pre><code class="lang-auto">with torch.no_grad():
    original_img_features = model(original_img) 
    style_features = model(style_img)
</code></pre>
<p>This will detach these two features from the computation graph, so it won’t be a problem if you compute loss using them multiple times.</p>