CapsuleNet RuntimeError

Alva-2020: <pre><code class="lang-auto">class DigitCapsLayer(nn.Module):
   def __init__(self, n_caps=120, n_dims=16, n_routing_iterations=3):
       super(DigitCapsLayer, self).__init__()


       self.n_caps = n_caps
       self.n_dims = n_dims
       self.n_routing_iterations = n_routing_iterations

       # init weight matrix W
       #self.routing_weights = nn.Parameter(torch.randn(1, 32*6*6, n_caps, n_dims, 8))
       self.routing_weights = nn.Parameter(torch.randn(1, 7744, n_caps, n_dims, 8)).cuda()
       #print(sum(p.numel() for p in self.routing_weights.parameters()) / 1e6)

   def squash(self, tensor=None, dim=-1, epsilon=1e-7):
       squared_norm = (tensor ** 2).sum(dim=dim, keepdim=True)
       safe_norm = torch.sqrt(squared_norm + epsilon)
       squash_factor = squared_norm / (1. + squared_norm)
       unit_vector = tensor / safe_norm
       squashed = squash_factor * unit_vector
       #return squash_factor * unit_vector
       return squashed


   def forward(self, x):
       # x: [Batch_size, 1152, 8]

       # create W and x
       batch_size = x.size(dim=0)

       W = self.routing_weights.repeat(batch_size, 1, 1, 1, 1) # [BS, 1152, 10, 16, 8]
       
       x = torch.unsqueeze(x, dim=-1).clone()
       x = torch.unsqueeze(x, dim=2).clone()
       x = x.repeat(1, 1, 120, 1, 1) # [BS, 1152, 10, 8, 1]

       # Compute the Predicted Output Vectors (u_hat[j|i])
       prediction_vectors = W @ x # [BS, 1152, 10, 16, 1]
       

       # Routing by agreement
      
       # Initial logits b[i, j]
       #initial_logits = Variable(torch.zeros(batch_size, 32*6*6, self.n_caps, 1, 1)).to(self.device) # [BS, 1152, 10, 1, 1]  initial_logits = Variable(torch.zeros(batch_size, 32*6*6, self.n_caps, 1, 1)).to(self.device) # [BS, 1152, 10, 1, 1]
       initial_logits = Variable(torch.zeros(batch_size, 7744, self.n_caps, 1, 1)).cuda() #.to(self.device) # [BS, 1152, 10, 1, 1]
       for i in range(self.n_routing_iterations):
           # The coupling coefficients c[i] = softmax(b[i])
           coupling_coefs = F.softmax(initial_logits, dim=2)
           
           outputs = (coupling_coefs * prediction_vectors).sum(dim=1, keepdims=True) # [BS, 1, 10, 16, 1]
           outputs = self.squash(tensor=outputs, dim=-2)

           if i != self.n_routing_iterations - 1:
               #outputs = outputs.repeat(1, 32*6*6, 1, 1, 1)
               outputs = outputs.repeat(1, 7744, 1, 1, 1)
               agreement = torch.transpose(prediction_vectors, -2, -1) @ outputs # [BS, 1152, 10, 1, 1]
               initial_logits += agreement 
       
       return outputs

class CapsNet(nn.Module):
   def __init__(self, n_routing_iterations=3, input_channel=1):
       super(CapsNet, self).__init__()


       # Layer 1: Convolution
       #self.Conv1 = nn.Conv2d(in_channels=3, out_channels=256, kernel_size=9, stride=1, padding=0)
       self.Conv1 = nn.Conv2d(in_channels=input_channel, out_channels=8, kernel_size=9, stride=1, padding=0)
       # Layer 2: Primary Capsules
       self.Conv2 = nn.Conv2d(in_channels=8, out_channels=16, kernel_size=9, stride=2, padding=0)

       # Layer 3: Digit Capsules
       self.digitcaps = DigitCapsLayer(n_routing_iterations=n_routing_iterations)
       # Reconstruction
       self.decoder = nn.Sequential(
           nn.Linear(16*120, 512),
           nn.ReLU(),
           nn.Linear(512, 1024),
           nn.ReLU(),
           nn.Linear(1024, 80*257),
           nn.Sigmoid()
       )

   def squash(self, tensor=None, dim=-1, epsilon=1e-7):
       squared_norm = (tensor ** 2).sum(dim=dim, keepdim=True)
       safe_norm = torch.sqrt(squared_norm + epsilon)
       squash_factor = squared_norm / (1. + squared_norm)
       unit_vector = tensor / safe_norm
       squashed = squash_factor * unit_vector
       return squashed
       #return squash_factor * unit_vector


   def safe_norm(self, tensor=None, dim=-1, epsilon=1e-7, keep_dims=False):
       squared_norm = (tensor ** 2).sum(dim=dim, keepdim=keep_dims)
       return torch.sqrt(squared_norm + epsilon)


   def forward(self, x, y=None):
       # 1st layer
       x = self.Conv1(x)
       x = F.relu(x, inplace=True)

       # primary capsules
       x = self.Conv2(x)
       x = F.relu(x, inplace=True)
       
       #强制重构矩阵维度
       #x = torch.reshape(x, (-1, 32*6*6, 8))
       x = torch.reshape(x, (x.size()[0],-1, 8))
       x = self.squash(tensor=x)
       '''
       caps_2_output----&gt; torch.Size([128, 1, 10, 16, 1])
       y_proba----&gt; torch.Size([128, 1, 10, 1])
       y_proba_argmax----&gt; torch.Size([128, 1, 1])
       y_pred--------&gt; torch.Size([128, 1])
       outputs---&gt;caps_2_outputs: torch.Size([128, 1, 10, 16, 1])
       outputs---&gt;y_pred: torch.Size([128])
       outputs---&gt;reconstructions: torch.Size([128, 784])
       '''
       # digit capsules
       caps_2_output = self.digitcaps(x) # [BS, 1, classNo., 16, 1]
       # reconstruction
       '''
       y_proba----&gt; torch.Size([8, 1, 120, 1])
       '''
       # find longest vector
       y_proba = self.safe_norm(caps_2_output, dim=-2)
       #y_proba----&gt; torch.Size([128, 1, 10, 1])
       y_proba_argmax = torch.argmax(y_proba, dim=2)
       #y_proba_argmax----&gt; torch.Size([128, 1, 1])
       y_pred = torch.squeeze(y_proba_argmax, dim=2)
       y_pred = torch.squeeze(y_pred, dim=1)
       mask = None
       if y is None:
           mask = torch.eye(120).cuda().index_select(dim=0, index=y_pred)
       else:
           mask = torch.eye(120).cuda().index_select(dim=0, index=y)
       

       reconstruction_mask_reshaped = torch.reshape(mask, [-1, 1, 120, 1, 1])
       
       caps2_output_masked = torch.multiply(caps_2_output, reconstruction_mask_reshaped)

       decoder_input = torch.reshape(caps2_output_masked, [-1, 120 * 16])

       reconstructions = self.decoder(decoder_input)

       #return caps_2_output, y_pred, reconstructions
       return caps_2_output, y_proba, reconstructions

if __name__ == '__main__':
   model = CapsNet(n_routing_iterations=3)
   print(sum(p.numel() for p in model.parameters()) / 1e6)
   x = torch.Tensor(8,1,257,80)
   '''
   (Pdb) p caps_2_output.size()
   torch.Size([8, 1, 120, 16, 1])
   (Pdb) p y_pred.size()
   torch.Size([8])
   (Pdb) p reconstructions.size()
   torch.Size([8, 20480]
   '''
   model(x)
``` when I use loss.backward() to trian the above CapNet， there is always an error "RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation: [torch.cuda.FloatTensor [8, 7744, 120, 1, 1]], which is output 0 of AddBackward0, is at version 3; expected version 2 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."  Please help me figure out which line occurs this error and how to fix them. Thanks so  much.</code></pre>

tom: <p>Have you looked at dropping the <code>inplace=True</code>?</p>
<p>Best regards</p>
<p>Thomas</p>