ca5ec9585e711f63acbe337bc38c37520189b9d7,se3_cnn/convolution.py,SE3KernelCombination,backward,#SE3KernelCombination#Any#,152

Before Change

                    for jj in range(self.multiplicites_in[j]):
                        sj = slice(begin_j + jj * self.dims_in[j], begin_j + (jj + 1) * self.dims_in[j])

                        grad_weight[weight_index : weight_index + b_el] = torch.mm(basis_kernels_ij, grad_kernel[si, sj].contiguous().view(-1, 1)).view(-1)
                        weight_index += b_el

                begin_j += self.multiplicites_in[j] * self.dims_in[j]
            begin_i += self.multiplicites_out[i] * self.dims_out[i]

After Change

            for j, mj in enumerate(self.multiplicites_in):
                b_el = self.kernels[i][j].size(0)
                basis_kernels_ij = self.kernels[i][j] // [beta, i, j, x, y, z]
                basis_kernels_ij = basis_kernels_ij.view(b_el, -1) // [beta, i*j*x*y*z]

                si = slice(begin_i, begin_i + mi * self.dims_out[i])
                sj = slice(begin_j, begin_j + mj * self.dims_in[j])

                grad = grad_kernel[si, sj] // [I * i, J * j, x, y, z]
                grad = grad.contiguous().view(mi, self.dims_out[i], mj, self.dims_in[j], -1).transpose(1, 2) // [I, J, i, j, x*y*z]
                grad = grad.contiguous().view(mi * mj, -1) // [I*J, i*j*x*y*z]
                grad = torch.mm(grad, basis_kernels_ij.transpose(0, 1)) // [I*J, beta]

                grad_weight[weight_index : weight_index + mi * mj * b_el] = grad.view(-1) // [I * J * beta]
                weight_index += mi * mj * b_el

                begin_j += self.multiplicites_in[j] * self.dims_in[j]

Instances