RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation about s2cnn

[mariogeiger]

I can fix it with torch.einsum("ij,jk->ik", (x.clone(), torch.randn(3, 3)))

[mariogeiger]

The problem comes from s2_rft when we use torch.einsum. The problem can be reproduced by the following code:

x = torch.randn(3, 3, requires_grad=True)
z1 = torch.einsum("ij,jk->ik", (x, torch.randn(3, 3)))
z2 = torch.einsum("ij,jk->ik", (x, torch.randn(3, 3)))
z1.sum().backward()

[mariogeiger]

Hi,
No I never observed this error, we always used a mono-model.
Did you try to simplify the model to see if the error still occur ? For instance using only s2conv or only so3conv ?

[udonuser]

Yes, I got same error using only s2conv in following code.

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np

from torch.utils.data import DataLoader
from torchvision import transforms
from torchvision.datasets import MNIST

from s2cnn import s2_near_identity_grid, S2Convolution


def S2conv2d(in_c, out_c, in_b, out_b):
    grid = s2_near_identity_grid(n_alpha=2 * in_b)
    return S2Convolution(in_c, out_c, in_b, out_b, grid)

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.conv1 = S2conv2d(1, 5, 14, 7)

    def forward(self, x):
        return self.conv1(x)


def main():
    use_cuda = torch.cuda.is_available()
    device = torch.device("cuda" if use_cuda else "cpu")
    
    WORKERS=1

    img_transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,)),
    ])

    train_loader = DataLoader(MNIST('./data', train=True, transform=img_transform, download=True),
                              batch_size=256, num_workers=WORKERS,  pin_memory=True, shuffle=True)

    model = Model().to(device)
    optimizer = torch.optim.SGD(model.parameters(), lr=1e-3,
                                momentum=0.9, weight_decay=5e-4)

    def train():
        model.train()
        for batch_idx, (image,target) in enumerate(train_loader):
            image = image.to(device)
            optimizer.zero_grad()

            # multi model
            output1 = model(image)
            output2 = model(image)
            loss = (output1 + output2).mean()

            loss.backward()
            optimizer.step()
            print("OK")
            break
    train()

[udonuser]

Thank you so much!