Inplace operation error in control problem

Question

I'm new to pytorch and I'm having a problem with some code to train aa neural network to solve a control problem. I use the following code to solve a toy version of my problem:

# SOME IMPORTS
import torch
import torch.autograd as autograd  
from torch import Tensor               
import torch.nn as nn                   
import torch.optim as optim       

# Device configuration
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')


# PARAMETERS OF THE PROBLEM
layers = [4, 32, 32, 4]  # Layers of the NN
steps = 10000  # Simulation steps
train_step = 1  # I train the NN for 1 epoch every train_step steps
lr = 1e-3  # Learning rate

After this I define a very simple network:

# DEFINITION OF THE NETWORK (A SIMPLE FEED FORWARD)

class FCN(nn.Module):
    
    def __init__(self,layers):
        super(FCN, self).__init__() #call __init__ from parent class 
         
        self.linears = []
        for i in range(len(layers)-2):
            self.linears.append(
                nn.Linear(layers[i], layers[i+1])
            )
            self.linears.append(
                nn.ReLU()
            )
            
        self.linears.append(
                nn.Linear(layers[-2], layers[-1])
        )
        
        self.linear_stack = nn.Sequential(*self.linears)
            
    'forward pass'
    def forward(self,x):
        
        out = self.linear_stack(x)

        return out

I then use the defined class to create my model:

model = FCN(layers)
model.to(device)
params = list(model.parameters())
optimizer = torch.optim.Adam(model.parameters(),lr=lr,amsgrad=False)

Then I define the loss function and the simulation function, ie the function that updates the state of my problem.

def simulate(state_old, model):

    state_new = model(state_old)

    return state_new

def lossNN(state_old,state_new, model):

    error = torch.sum( (state_old-state_new)**2 )

    return error

And finally I train my model:

torch.autograd.set_detect_anomaly(True)

state_old = torch.Tensor([0.01, 0.01, 0.5, 0.1]).to(device)

for i in range(steps):
    state_new = simulate(state_old, model)
    
    if i%train_step == 0:
        optimizer.zero_grad()
        loss = lossNN(state_old, state_new, model)
        loss.backward(retain_graph=True)
        optimizer.step()
        
    state_old = state_new
    
    
    if (i%1000)==0:
        print(loss)
        print(state_new)

I then get the following error. Here you can find the backtrace:

RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation: [torch.cuda.FloatTensor [32, 4]], which is output 0 of AsStridedBackward0, is at version 2; expected version 1 instead. Hint: the backtrace further above shows the operation that failed to compute its gradient. The variable in question was changed in there or anywhere later. Good luck!

Answer 1

You need to use detach to remove the gradient created in the previous state.

state_old = state_new

state_old = state_new.detach()

Then your training code changes to:

torch.autograd.set_detect_anomaly(True)

state_old = torch.Tensor([0.01, 0.01, 0.5, 0.1]).to(device)

for i in range(steps):
    state_new = simulate(state_old, model)

    if i%train_step == 0:
        optimizer.zero_grad()
        loss = lossNN(state_old, state_new, model)
        loss.backward(retain_graph=True)
        optimizer.step()

    state_old = state_new.detach()
    
    if (i%1000)==0:
        print(loss)
        print(state_new)