Torch.gather: uncontrollable randomness (quirky)

Yochengliu: <p><strong>[Problem description]</strong> I use <em>torch.gather</em> to index a tensor, and find uncontrollable randomness. Example code is below:</p>
<pre><code class="lang-auto"># batchsize=16, feature dim=3, pixels = 128
feature = torch.randn(16, 3, 128) 
# index feature: find 64 pixels in 128 and search 24 relevant pixels all over 64 pixels.
pixel_index = torch.randint(0, 128, (16, 64, 24)).long().cuda()
pixel_index = pixel_index.view(16, 64*24).unsqueeze(1).expand(-1, 3, -1)  # (16, 3, 64*24)
# get the indexed feature with shape (16, 3, 64, 24)
indexed_feature = torch.gather(feature, dim=2, index=pixel_index).view(16, 3, 64, 24)
</code></pre>
<p>Only using <em>torch.gather</em> <strong>once</strong> can get determined result while using it <strong>above twice</strong> would lead to undetermined result. <strong>All random seeds are fixed</strong>. I have found the key issue locates in the backward of <em>torch.gather</em>, but <em>why? Is there a way to fix this quirky randomness?</em></p>
<p><strong>[Test code for reproducing the issue]</strong></p>
<pre><code class="lang-auto">import torch
import torch.optim as optim
import torch.nn as nn
import numpy as np
import os
import random

seed = 123
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)            
torch.cuda.manual_seed(seed)       
torch.cuda.manual_seed_all(seed) 
os.environ['PYTHONHASHSEED'] =str(seed)
torch.backends.cudnn.benchmark = False
torch.backends.cudnn.deterministic = True

def gather_op(feat, idx):
    """
    features : (B, C, N)
    idx : (B, M)
    return: (B, C, M)
    """
    return torch.gather(feat, 2, idx.unsqueeze(1).expand(-1, feat.shape[1], -1))

class sampling(nn.Module):
    def __init__(self):
        super().__init__()
    def forward(self, x, idx):
        B, N, M = idx.shape
        return gather_op(x, idx.view(B, N*M)).view(B, -1, N, M)
    
class model_test1(nn.Module):    # only use gather_op once
    def __init__(self, num_classes, input_dim=3):
        super().__init__()
       
        self.conv1d_L1 = nn.Conv1d(input_dim, 32, 1)
        self.conv2d_L2 = nn.Conv2d(32, 128, 1)
        self.relu = nn.ReLU(inplace=True)
        
        self.sampling = sampling()
        
        self.classifier = nn.Linear(128, 10)
    def forward(self, x, idx):
        x = self.relu(self.conv1d_L1(x))      # (16, 32, 128)
        x = self.sampling(x, idx)             # (16, 32, 64, 24)
        x = self.relu(self.conv2d_L2(x))      # (16, 128, 64, 24)
        x = torch.sum(torch.sum(x, dim = 3), dim = 2) # (16, 128)
        x = self.classifier(x)      # (16, 10)
        return x
    
class model_test2(nn.Module):    # use gather_op twice
    def __init__(self, num_classes, input_dim=3):
        super().__init__()
       
        self.conv1d_L1 = nn.Conv1d(input_dim, 32, 1)
        self.conv2d_L2 = nn.Conv2d(32, 128, 1)
        self.conv2d_L3 = nn.Conv2d(128, 256, 1)
        self.relu = nn.ReLU(inplace=True)
        
        self.sampling = sampling()
        
        self.classifier = nn.Linear(256, 10)
    def forward(self, x, idx1, idx2):
        x = self.relu(self.conv1d_L1(x))      # (16, 32, 128)
        x = self.sampling(x, idx1)            # (16, 32, 64, 24)
        x = self.relu(self.conv2d_L2(x))      # (16, 128, 64, 24)
        x = torch.sum(x, dim = 3)             # (16, 128, 64)
        x = self.sampling(x, idx2)            # (16, 128, 32, 24)
        x = self.relu(self.conv2d_L3(x))      # (16, 256, 32, 24)
        x = torch.sum(torch.sum(x, dim = 3), dim = 2) # (16, 256)
        x = self.classifier(x)      # (16, 10)
        return x
    

model1, model2 = model_test1(10, 3), model_test2(10, 3)
model1, model2 = model1.cuda(), model2.cuda()
optimizer1, optimizer2 = optim.Adam(model1.parameters(), lr=0.001), optim.Adam(model2.parameters(), lr=0.001)
    
data = torch.randn(1000, 3, 128).cuda()  # 1000 samples, 3-d features, 128 pixels
label = torch.randint(0, 10, (1000,)).long().cuda()

criterion = nn.CrossEntropyLoss()
model1.train()
model2.train()
for iter in range(60):
    index = torch.arange(16*iter, 16*(iter+1)).long().cuda()  # batchsize = 16
    input, target = data[index, ...], label[index, ...]
    optimizer1.zero_grad()
    optimizer2.zero_grad()
    
    # index to sample 24 pixels
    idx1 = torch.randint(0, 128, (16, 64, 24)).long().cuda()  # (16, 64, 24)
    idx2 = torch.randint(0, 64, (16, 32, 24)).long().cuda()   # (16, 32, 24)
    pred1, pred2 = model1(input, idx1), model2(input, idx1, idx2)
    target = target.view(-1)
    loss1, loss2 = criterion(pred1, target), criterion(pred2, target)
    loss1.backward()
    loss2.backward()
    optimizer1.step()
    optimizer2.step()
    
    print('[%2d] loss1: %0.16f \t [%2d] loss2: %0.16f' % (iter, loss1.data.clone(), iter, loss2.data.clone()))
</code></pre>
<p><strong>[Results show]</strong><br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/6/0/60bb55487954bbeb527f7ed27a95376dbf90977b.png" data-download-href="https://discuss.pytorch.org/uploads/default/60bb55487954bbeb527f7ed27a95376dbf90977b" title="1575981565(1)"><img src="https://discuss.pytorch.org/uploads/default/original/3X/6/0/60bb55487954bbeb527f7ed27a95376dbf90977b.png" alt="1575981565(1)" data-base62-sha1="dNJcbzlTn6XPFm3f4y6dAER6e2D" width="690" height="211" data-dominant-color="535351"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">1575981565(1)</span><span class="informations">1553×476 35.8 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div><br>
<em>loss1</em> of model_test1 that uses <em>torch.gather</em> once is determined.<br>
<em>loss2</em> of model_test2 that uses <em>torch.gather</em> twice is undetermined in the two runs, as highlighted in the figure.</p>
<p><strong>[Environment]</strong><br>
14.04.1-Ubuntu<br>
cuda8.0 + cudnn v7.1<br>
anaconda3 (Python 3.6.5)<br>
pytorch 0.4.0/1.0.0 (both have this problem)<br>
numpy 1.14.3</p>
<p><strong>[Why not try other indexing methods]</strong><br>
I have tried the two methods below (for some reasons, I prefer to use <em>pytorch 0.4.0</em> in my project):</p>
<p><em>METHOD 1: directly use advanced indexing</em></p>
<pre><code class="lang-auto">B, N, M = idx.shape
batch_indices = torch.arange(B).type(torch.cuda.LongTensor).view(B, 1, 1).expand(-1, N, M)
return feat.transpose(1, 2)[batch_indices, idx, :].permute(0, 3, 1, 2)   # (B, C, N, M)
</code></pre>
<p>No randomness indeed, but in <em>pytorch 0.4.0</em>, this will cause huge memory issue and is very slow. Moreover, directly indexing for each batch (<em>for i in range(batchsize)</em>) also suffers the same memory issue.</p>
<p><em>METHOD 2: use torch.index_select</em></p>
<pre><code class="lang-auto">B = idx.shape[0]
new_feat = [torch.index_select(feat[i, ...], 1, idx[i, ...]).unsqueeze(0) for i in range(B)]
return torch.cat(new_feat, 0)
</code></pre>
<p>There is also randomness in <em>torch.index_select</em>.</p>
<p><em>torch.gather</em> for me is fast and memory-saving, thus I want to use it for indexing.</p>
<p><strong>Any useful help is welcome! Thanks.</strong></p>

albanD: <p>Hi,</p>
<p>Have you read the documentation note about <a href="https://pytorch.org/docs/stable/notes/randomness.html">reproducibility</a> ?<br>
I’m afraid this is the kind of issue you’re seing here. Both reasults are correct wrt floating point computations but the gradient descent you’re doing is making any small difference significant.</p>

Yochengliu: <p><a class="mention" href="/u/alband">@albanD</a>  Thanks. I have read it. So do you mean the backward of <em>torch.gather</em> uses some operations like <em>atomicAdd</em> that are undetermined anyway? If so, why using <em>torch.gather</em> once can get determined results?</p>

albanD: <blockquote>
<p>If so, why using  <em>torch.gather</em>  once can get determined results?</p>
</blockquote>
<p>I don’t know for sure. But it might be that with your particular hardware, a single gather happen to most of the time return bit-wise exact results. While queuing more than one expose a different behavior because it puts the gpu in a different state than it was when you call only one?</p>
<p>Why do you need exact determinism? knowing that it won’t work as soon as you change machine/pytorch version/dependency version?</p>

Yochengliu: <p><a class="mention" href="/u/alband">@albanD</a>  Maybe. I can’t reproduce the performance of my model if there is any randomness (the gap is non-negligible). Until now, I haven’t found a good solution for this issue. Thanks anyway.</p>

albanD: <p>Hi,</p>
<p>Does that mean that your model is very unstable and only with one random seed, you get good performance by chance?</p>

Yochengliu: <p><a class="mention" href="/u/alband">@albanD</a>  Yes. With the uncontrollable randomness in <em>torch.gather</em>, I can’t even distinguish whether my ideas lead to the improvement of performance, or it is just by chance. This really troubles me…</p>

albanD: <p>My personal opinion is that if you floating point precision makes your model not work. It is not good. Because no one will ever be able to reproduce your result. So I would conclude that this particular seed worked by chance.</p>

Yochengliu: <p><a class="mention" href="/u/alband">@albanD</a>  Thanks.  I’ll check my model carefully.</p>