Outer Loop Backward with Inner Loop no_grad

[Fat-Fat-Lee]

Outer Loop Backward with Inner Loop no_grad

Question

I encountered an issue where I defined a tensor self.biggest_mask inside the inner loop of ImplicitModel, which is calculated by a custom function save_biggest_mask. Because save_biggest_mask involves some for loops and operations that do not support gradient computation, I do not want this tensor to participate in gradient propagation.

However, in the outer loop, torchopt still performed gradient propagation on this tensor, leading to RuntimeError: Cannot access data pointer of Tensor that doesn't have storage.

I have tried using torch.no_grad() and torchopt.stop_gradient(self.biggest_mask) ,but don't work.

Inner Loop Code

def get_implicit_model(solver):
    class ImplicitModel(torchopt.nn.ImplicitMetaGradientModule,
                        linear_solve=solver):
        def __init__(self, clser, atter, lamb, max_epochs=200, optimizer='Adam', lr=0.01,device=None,biggest_mask=None):
            super().__init__()
            self.register_meta_module('atter', atter)
            object.__setattr__(self, 'clser', clser)
            object.__setattr__(self, 'biggest_mask',biggest_mask)
            self.lr = lr
            self.lamb = lamb
            self.max_epochs = max_epochs
            self.optimizer = optimizer
            self.device=device

        def reset(self, z):
            self.z0 = z
            self.z = nn.Parameter(z.clone().detach_(), requires_grad=True)


        def objective(self):
            mask_ = self.atter(self.z)
            

            #------------todo biggest_mask is here------------------------#
            with torch.no_grad():
                self.biggest_mask=save_biggest_mask(mask_.detach().clone(), self.device)
                torchopt.stop_gradient(self.biggest_mask)
            mask_ = mask_ * self.biggest_mask
            #--------------------------------------------------------------#
            
            
            pred = self.clser(mask_ * self.z)
            celoss = malign_celoss(pred, self.target)
            norm = torch.norm(self.z - self.z0)
            return celoss + self.lamb * norm

        @torch.enable_grad()
        def solve(self, ):
            optimizer = getattr(torch.optim, self.optimizer)(params=[self.z], lr=self.lr)
            mask_=self.atter(self.z0)
            pred=self.clser(mask_* self.z0)
            mscore=nn.Softmax(dim=1)(pred)[:, 1]
            self.target = 1 - (mscore >= 0.5).long().clone().detach()
            for epoch in range(self.max_epochs):
                optimizer.zero_grad()
                
                #------------todo biggest_mask is here------------------------#
                torchopt.stop_gradient(self.biggest_mask)
                loss = self.objective()
                #-------------------------------------------------------------#
                
                norm = torch.norm(self.z - self.z0)
                celoss = loss - self.lamb * norm
                loss.backward(inputs=[self.z])
                optimizer.step()
              
            return 

    return ImplicitModel

Outter Loop Code

            for batch_idx, (train_z, train_mask,_) in enumerate(train_loader):
                rand_flip_(train_z, train_mask)
                train_z, train_mask= train_z.to(device), train_mask.to(device)
                optim_atter.zero_grad()

                imp.reset(train_z)
                
                #------------todo biggest_mask is here------------------------#
                inner_log = imp.solve()#todo imp inner loop
                torchopt.stop_gradient(imp.biggest_mask)
                #-------------------------------------------------------------#

                diff = imp.z - train_z
                reward = inside_nod_loss(diff, train_mask)*args.ratio
                punish = beyond_nod_loss(diff, train_mask)
                alignloss = punish - reward
                alignloss.backward()#todo outer loop loss backward
                optim_atter.step()

save_biggest_mask

def save_biggest_mask(A,device):
    with torch.no_grad():
        A_ = digitize(minmax_normalize(A, (2, 3))).to(torch.uint8).to(device)
        B = torch.full_like(A_, fill_value=0.1, dtype=torch.float32).to(device)
        A_ = cc_torch.connected_components_labeling(A_.squeeze(1)).unsqueeze(1)
        for i in range(A_.size(0)):
            A_i = A_[i, 0, :, :]
            B_i = B[i, 0, :, :]
            unique_values, counts = torch.unique(A_i, return_counts=True)
            counts[torch.where(unique_values == 0)] = 0
            biggest_value = unique_values[torch.argmax(counts).item()]
            B_i[torch.where(A_i == biggest_value)] = 1
    return B

error message

Traceback (most recent call last):
  File "/home/lijingwen/Projects/Counter_align/DDSM_hierarchical/hierachical_model/img2label/main_cross_avg3.py", line 161, in <module>
    alignloss.backward()
  File "/home/lijingwen/Miniconda3/envs/toy/lib/python3.8/site-packages/torch/_tensor.py", line 487, in backward
    torch.autograd.backward(
  File "/home/lijingwen/Miniconda3/envs/toy/lib/python3.8/site-packages/torch/autograd/__init__.py", line 197, in backward
    Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
  File "/home/lijingwen/Miniconda3/envs/toy/lib/python3.8/site-packages/torch/autograd/function.py", line 267, in apply
    return user_fn(self, *args)
  File "/home/lijingwen/Miniconda3/envs/toy/lib/python3.8/site-packages/torchopt/diff/implicit/decorator.py", line 341, in backward
    vjps = _root_vjp(
  File "/home/lijingwen/Miniconda3/envs/toy/lib/python3.8/site-packages/torchopt/diff/implicit/decorator.py", line 106, in _root_vjp
    _, optimality_cond_vjp_fn, *_ = functorch.vjp(optimality_cond, solution)
  File "/home/lijingwen/Miniconda3/envs/toy/lib/python3.8/site-packages/functorch/_src/eager_transforms.py", line 262, in vjp
    return _vjp_with_argnums(func, *primals, has_aux=has_aux)
  File "/home/lijingwen/Miniconda3/envs/toy/lib/python3.8/site-packages/functorch/_src/vmap.py", line 35, in fn
    return f(*args, **kwargs)
  File "/home/lijingwen/Miniconda3/envs/toy/lib/python3.8/site-packages/functorch/_src/eager_transforms.py", line 289, in _vjp_with_argnums
    primals_out = func(*primals)
  File "/home/lijingwen/Miniconda3/envs/toy/lib/python3.8/site-packages/torchopt/diff/implicit/decorator.py", line 104, in optimality_cond
    return optimality_fn(solution, *args)
  File "/home/lijingwen/Miniconda3/envs/toy/lib/python3.8/site-packages/torchopt/diff/implicit/nn/module.py", line 73, in _stateless_optimality_fn
    return self.optimality(*input, **kwargs)
  File "/home/lijingwen/Miniconda3/envs/toy/lib/python3.8/site-packages/torchopt/diff/implicit/nn/module.py", line 91, in optimality
    flat_grads = objective_grad_fn(
  File "/home/lijingwen/Miniconda3/envs/toy/lib/python3.8/site-packages/functorch/_src/eager_transforms.py", line 1241, in wrapper
    results = grad_and_value(func, argnums, has_aux=has_aux)(*args, **kwargs)
  File "/home/lijingwen/Miniconda3/envs/toy/lib/python3.8/site-packages/functorch/_src/vmap.py", line 35, in fn
    return f(*args, **kwargs)
  File "/home/lijingwen/Miniconda3/envs/toy/lib/python3.8/site-packages/functorch/_src/eager_transforms.py", line 1111, in wrapper
    output = func(*args, **kwargs)
  File "/home/lijingwen/Miniconda3/envs/toy/lib/python3.8/site-packages/torchopt/diff/implicit/nn/module.py", line 54, in _stateless_objective_fn
    return self.objective(*input, **kwargs)
  File "/home/lijingwen/Projects/Counter_align/DDSM_hierarchical/hierachical_model/model.py", line 41, in objective
    self.biggest_mask=save_biggest_mask(mask_.detach().clone(), self.device)
  File "/home/lijingwen/Projects/Counter_align/DDSM_hierarchical/hierachical_model/tools/util.py", line 177, in save_biggest_mask
    A_ = cc_torch.connected_components_labeling(A_.squeeze(1)).unsqueeze(1)
  File "/home/lijingwen/Miniconda3/envs/toy/lib/python3.8/site-packages/cc_torch-0.1-py3.8-linux-x86_64.egg/cc_torch/connected_components.py", line 18, in connected_components_labeling
    return _C.cc_3d(x)
RuntimeError: Cannot access data pointer of Tensor that doesn't have storage

Process finished with exit code 1

@XuehaiPan @Benjamin-eecs @JieRen98 @waterhorse1

[XuehaiPan]

It seems that the biggest_mask attribute is data-dependent. You should implement your methods as pure-function. There should not be data assignment (e.g. tensor[idx] = value) or .data access (e.g., tensor.data or tensor.numpy()) in your code.

See also:

• [Bug/functorch] Cannot use tensor.detach().numpy() for GradTrackingTensor: Cannot access data pointer of Tensor that doesn't have storage pytorch/pytorch#91810

[Fat-Fat-Lee]

Thank you very much for XuehaiPan's response, but I still haven't solved the problem.

1. In my save_biggest_mask implementation, only A_ = A.clone() is operated, and the input A is not operated on, but I still get an error.This should be considered a pure function. However, I still encounter errors.
2. After using with torch.no_grad(): in save_biggest_mask, the output results can be treated as tensors without gradients and can be normally backwarded in a regular network. Why doesn't it work with torchopt?
3. I tried modifying save_biggest_mask based on the content in the link provided by XuehaiPan ([Bug/functorch] Cannot use tensor.detach().numpy() for GradTrackingTensor: Cannot access data pointer of Tensor that doesn't have storage pytorch/pytorch#91810), and the error return torch.TypedStorage(wrap_storage=self._storage(), dtype=self.dtype) NotImplementedError: Cannot access storage of TensorWrapper still remains. Moreover, I believe this also indicates that torchopt still does not treat the result of the save_biggest_mask function as a tensor without gradients, and it still attempts to differentiate the tensor internally.

def detach_numpy(tensor):
    tensor = tensor.detach().cpu()
    if torch._C._functorch.is_gradtrackingtensor(tensor):
        tensor = torch._C._functorch.get_unwrapped(tensor)
        return np.array(tensor.storage().tolist()).reshape(tensor.shape)
    return tensor.numpy()

def save_biggest_mask(A,device):
    with torch.no_grad():

        if A.size(0) % 2 == 1:
            A_ = torch.empty(A.size(0) + 1, A.size(1),A.size(2),A.size(3))
            A_[:A.size(0)] = A.clone()
            A_[A.size(0)] = A[A.size(0) - 1].clone()
        else:
            A_ = A.clone()

        A_ = digitize(minmax_normalize(A_, (2, 3))).to(torch.uint8).to(device)
        B = torch.full_like(A, fill_value=0.1, dtype=torch.float32).to(device)

        A_ = torch.tensor(detach_numpy(A_)).to(device)  # todo
        A_ = cc_torch.connected_components_labeling(A_.squeeze(1)).unsqueeze(1)

        for i in range(A.size(0)):
            A_i = A_[i, 0, :, :]
            B_i = B[i, 0, :, :]
            unique_values, counts = torch.unique(A_i, return_counts=True)
            counts[torch.where(unique_values == 0)] = 0
            biggest_value = unique_values[torch.argmax(counts).item()]
            B_i[torch.where(A_i == biggest_value)] = 1
        return B

What should I do? How should I modify my code? @XuehaiPan @Benjamin-eecs @waterhorse1 @JieRen98

[Fat-Fat-Lee]

move the save_biggest_mask function out of self.objective,put it into self.solve instead~