Sobel loss per channel of multiple gray image like

username: David_Sriker1
cooked: <p>I have an input model which is a bayer image → (b, 4, h, w)<br>
I want to apply sobel loss to each channel (treat it as if it was gray image):</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
import torch.nn.functional as F

from .data_record import DataRecord


class GradLayer(nn.Module):
    def __init__(self):
        super(GradLayer, self).__init__()
        kernel_v = [[1, 0, -1],
                    [2, 0, -2],
                    [1, 0, -1]]
        kernel_h = [[1, 2, 1],
                    [0, 0, 0],
                    [-1, -2, -1]]
        kernel_h = torch.FloatTensor(kernel_h).unsqueeze(0).unsqueeze(0)
        kernel_v = torch.FloatTensor(kernel_v).unsqueeze(0).unsqueeze(0)
        self.weight_h = nn.Parameter(data=kernel_h, requires_grad=False)
        self.weight_v = nn.Parameter(data=kernel_v, requires_grad=False)

    @staticmethod
    def get_gray(x):
        '''
        Convert image to its gray one.
        '''
        gray_coeffs = [65.738, 129.057, 25.064]
        convert = x.new_tensor(gray_coeffs).view(1, 3, 1, 1) / 256
        x_gray = x.mul(convert).sum(dim=1)
        return x_gray.unsqueeze(1)

    def forward(self, x):
        if x.shape[1] == 3:
            x = self.get_gray(x)

        _n_groups = x.shape[1]
        x_v = F.conv2d(x, self.weight_v.repeat((_n_groups, 1, 1, 1)), padding=1, groups=_n_groups)
        x_h = F.conv2d(x, self.weight_h.repeat((_n_groups, 1, 1, 1)), padding=1, groups=_n_groups)
        x = torch.sqrt(torch.pow(x_v, 2) + torch.pow(x_h, 2) + 1e-6)

        return x


class SobelLoss(nn.Module):
    def __init__(self):
        super(SobelLoss, self).__init__()
        self.loss = nn.L1Loss()
        self.grad_layer = GradLayer()
        self.unshuffle = nn.PixelUnshuffle(2)

    def forward(self, loss_input: DataRecord) -&gt; torch.Tensor:
        _target_seq = loss_input.target_sequence
        target = self.unshuffle(_target_seq[:, -1]) if len(_target_seq.shape) == 5 else _target_seq
        prediction = self.unshuffle(loss_input.prediction) if len(_target_seq.shape) == 5 else loss_input.prediction

        output_grad = self.grad_layer(prediction)
        gt_grad = self.grad_layer(target)
        return self.loss(output_grad, gt_grad)
</code></pre>
<p>Now I’m using group convolution but I’m getting confused on the group vs replicate the kernel:</p>
<p>group: <code>x_v = F.conv2d(x, self.weight_v.repeat((_n_groups, 1, 1, 1)), padding=1, groups=_n_groups)</code><br>
replicate: <code>x_v = F.conv2d(x, self.weight_v.repeat((_n_groups, _n_groups, 1, 1)), padding=1)</code></p>
<p>Can you please help elaborate on the matter and maybe give some ituition to the difference.<br>
In both I get the desired output shape (b, 4, h, w) but I wonder if it really the gradients or some unknown I’m missing</p>
<p>Thanks</p>

username: eqy
cooked: <p>It depends on what you are trying to do. As you observed, both give you the expected output shape, but the difference is that the first one will treat each channel independently, as you have <code>_n_groups</code> equal to the number of channels. That is, for each output channel, the result only depends on a <em>single</em> input channel—<code>x_v[:,0,:,:]</code> only depends on <code>x[:,0,:,:]</code>, and <code>x_v[:,1,:,:]</code> only depends on <code>x[:,1,:,:]</code> and so on.</p>
<p>For the second case, you would add together the the result across all of the channels as the number of groups is 1 by default. So now <code>x_v[:,0,:,:]</code> is computed using <code>x[:,0,:,:]</code> and <code>x[:,1,:,:]</code> and <code>x[:,2,:,:]</code> and so on. In fact, since your filters are the same for each <em>output</em> channel in this case, you should be able to verify that <code>x_v[:,0,:,:] == x_v[:,1,:,:] == x[:,2,:,:] == x[:,3,:,:]</code> (barring possible minor numerical differences).</p>
<p>If we call <code>x_v1</code> <code>x_v</code> when it is computed according to the first method (<code>_n_groups = x.shape[1]</code>) and <code>x_v2</code> when computed according to the second method, then <code>x_v2[:,0,:,:] == x_v1[:,0,:,:] + x_v1[:,1,:,:] + x_v1[:,2,:,:] + x_v1[:,3,:,:]</code>. Note that these two equations are a consequence of the filter repeats but not true in the general case where the filters are allowed to have arbitrary values.</p>
<p>As it seems strange that you would want to have the same values across all channels, I assume you would want to do the former, but I don’t know the context of your use-case.</p>