Custom CNN with softmax applied to each conv operation

username: tkelestemur
cooked: <p>I am trying to write a custom CNN layer that applies softmax to each convolution operation. So each pixel in the output image is gonna be valued between [0, 1] and it is the sum of the convolved pixel. An example of TensorFlow implementation can be seen <a href="https://github.com/AdaCompNUS/qmdp-net/blob/master/qmdpnet.py#L312" rel="nofollow noopener">here</a>.</p>
<p>Ideally, this should be trained with binary cross-entropy loss. I tried below but it does not train.</p>
<pre><code class="lang-auto">class TransitionModel(nn.Module):
    def __init__(self):
        super(TransitionModel, self).__init__()

        self.weight = nn.Parameter(torch.Tensor(1, 1, 3, 3))
        nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
        self.softmax = nn.Softmax(dim=1)

    def forward(self, s, test=False):
        return F.conv2d(s, self.softmax(self.weight), padding=1)

model = TransitionModel()
pred_s_p = model(s)
pred_s_p = pred_s_p.squeeze()
loss = F.binary_cross_entropy(pred_s_p, s_p)
</code></pre>
<p>s.shape = [1, 1, 32, 32]<br>
pred_s_p.shape = [32, 32]</p>
<p>It throws the following error:</p>
<pre><code class="lang-auto">RuntimeError: reduce failed to synchronize: cudaErrorAssert: device-side assert triggered
</code></pre>

username: ptrblck
cooked: <p>Could you check, if your code runs on the CPU?<br>
I assume <code>F.binary_cross_entropy</code> might throw an error, as it expects probabilities as the model output, while you are passing the raw output as <code>pred_s_p</code>.<br>
Try to use <code>F.binary_cross_entropy_with_logits</code> instead.</p>

username: tkelestemur
cooked: <p>The model and data are on GPU. <code>F.binary_cross_entropy_with_logits</code> works but the loss does not change. Since I am applying softmax, the values should be probabilities in the output that’s why I thought I should use <code>F.binary_cross_entropy</code>.</p>

username: ptrblck
cooked: <p>You are applying the softmax on the weights, not the output.<br>
Depending on the distribution of your input you will not get probabilities as the output, which would raise an error as:</p>
<pre><code class="lang-python">RuntimeError: Assertion `x &gt;= 0. &amp;&amp; x &lt;= 1.' failed. input value should be between 0~1, but got -1.429985
</code></pre>

username: tkelestemur
cooked: <p>I see, any idea how I can apply softmax to the output of each conv operation?</p>
<p>Thanks for quick responses btw.</p>

username: ptrblck
cooked: <p>You could just call it directly of the output of <code>F.conv2d</code>:</p>
<pre><code class="lang-python">return self.softmax(F.conv2d(s, self.weight, padding=1))
</code></pre>

username: tkelestemur
cooked: <p>Doesn’t this apply the softmax over all the pixels in the output image? I need to apply it to each conv operation. Let’s say we have a <code>kernel</code> with size of [3, 3] and <code>image</code> [10, 10].</p>
<p>The layer I want should do:</p>
<pre><code class="lang-auto">softmax(image[0:3, 0:3] * kernel), softmax(image[0:3, 1:4] * kernel) ... 
softmax(image[1:4, 0:3] * kernel), softmax(image[1:4, 1:4] * kernel) ...  
...
</code></pre>
<p>Does this makes sense?</p>

username: richardbaxter
cooked: <p>Is there a method to apply a softmax independently (parallelised) across multiple sets of inputs?</p>

username: richardbaxter
cooked: <p>A solution is to artificially extend the batch with the multiple inputs;</p>
<pre><code class="lang-auto">#Assume x of shape [numberOfSamples, numberOfSublayers, numberDimensions]
numberOfSamples = x.shape[0]
numberOfSublayers = x.shape[1]
x = pt.reshape(x, (x.shape[0]*x.shape[1], x.shape[2]))
x = self.softmax(x)
x = pt.reshape(x, (numberOfSamples, numberOfSublayers, x.shape[1]))
</code></pre>