Move the loss function to GPU

username: Jindong
cooked: <p>Hi, every one,</p>
<p>I have a question about the “.cuda()”. In an example of Pytorch, I saw that there were the code like this:</p>
<p>criterion = nn.CrossEntropyLoss().cuda()</p>
<p>In my code, I don’t do this. So I am wondering if it necessary to move the loss function to the GPU.</p>
<p>Thanks</p>

username: royboy
cooked: <p>If your input tensor is a cuda tensor, it will run the cuda loss function.</p>

username: ptrblck
cooked: <p>Additionally to what <a class="mention" href="/u/royboy">@royboy</a> said, you need to push your criterion to the GPU, if it’s stateful, i.e. if it has some parameters or internal states.<br>
Usually loss functions are just functional so that it is not necessary.</p>

username: Oscar_Rangel
cooked: <p>criterion = nn.CrossEntropyLoss().cuda() if torch.cuda.is_available() else nn.CrossEntropyLoss()</p>

username: samra-irshad
cooked: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> Could you explain what do you mean by ‘if the criterion is stateful or if it has some parameters or some internal states’? I  am not sure what does internal states mean</p>

username: ptrblck
cooked: <p>A <code>weight</code> parameter could be seen as an internal state and would yield a device mismatch error.<br>
Of course you might define the <code>weight</code> parameter as a CUDATensor, but you could also move the criterion to the device:</p>
<pre><code class="lang-python">output = torch.randn(10, 10, requires_grad=True, device='cuda')
target = torch.randint(0, 10, (10,), device='cuda')

weight = torch.empty(10).uniform_(0, 1)
criterion = nn.CrossEntropyLoss(weight=weight)

loss = criterion(output, target) # error
&gt; RuntimeError: Expected object of device type cuda but got device type cpu for argument #3 'weight' in call to _thnn_nll_loss_forward

criterion.cuda()
loss = criterion(output, target) # works
</code></pre>

username: samra-irshad
cooked: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> Right, I do not understand why did you give the weight parameter as an input to loss function in <code>criterion = nn.CrossEntropyLoss(weight=weight)</code> ? I have never seen anyone feeding loss function with weight.</p>

username: ptrblck
cooked: <p>The <code>weight</code> argument can be used to create a class weighting, as described in the <a href="https://pytorch.org/docs/stable/generated/torch.nn.CrossEntropyLoss.html#torch.nn.CrossEntropyLoss">docs</a> of the criterion. It’s sometimes used to e.g. counter overfitting effects of training a model on imbalanced datasets.<br>
Weighted loss functions are not new in deep learning and were already used in the “classical” machine learning domain.</p>

username: samra-irshad
cooked: <p>Right. So these are class weights. I am not sure if I correctly understand the meaning of “internal state”. How would you define it and does internal state varies from module to module? In source code for <code>nn.Softmax()</code>, I can see there are some lines on ‘state’, i.e.,</p>
<pre><code class="lang-auto">def __setstate__(self, state):
        self.__dict__.update(state)
        if not hasattr(self, 'dim'):
            self.dim = None
</code></pre>
<p>Is it the internal state you are refering towards? I also noticed the modules and functions inheriting <code>nn.Module</code> are mostly moved to CUDA, any thoughts?</p>

username: ptrblck
cooked: <p>By “internal states” I mean all class attributes in the modules.<br>
In particular buffers and parameters are of interest, as they would need to be pushed to the appropriate device (e.g. such as the <code>weight</code> buffer in <code>nn.CrossEntropyLoss</code>).<br>
The <code>dim</code> argument in <code>nn.Softmax</code> however will not be pushed to the device, as it’s a plain Python integer to specify the dimension the softmax is applied in.</p>

username: Mohamed_Desouky
cooked: <p>trying this and got no error</p>
<pre><code class="lang-auto">input = torch.randn(10, 10, requires_grad=True, device='cuda')

criterion = nn.Softmax(dim = 1)

loss = criterion(input) # no error
</code></pre>