Variable assignment on multiple GPUs

michaelklachko: <p>I need to set a boolean flag in the code running on multiple GPUs. When using a single GPU, <code>self.calculate_running</code> is being set to False correctly after the first iteration. It’s not being set when I use more than one GPU:</p>
<pre><code class="lang-auto">class PCTL_Layer(nn.Module):
    def __init__(self, calculate_running=False):
        super(PCTL_Layer, self).__init__()
        self.register_buffer('running', torch.zeros(1))
        self.calculate_running = calculate_running

    def forward(self, input):
        if self.calculate_running:
            pctl, _ = torch.kthvalue(input.view(-1), int(input.numel() * 0.9))
            self.running = pctl
            print(' gpu {}  calculate_running: {}'.format(torch.cuda.current_device(), self.calculate_running))
            self.calculate_running = False

        return self.running


class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.percentile = PCTL_Layer(calculate_running=True)

    def forward(self, x):
        x = self.percentile(x)
        return x

model = Model()
model = torch.nn.DataParallel(model).cuda()

for i, image in enumerate(train_loader):
    output = model(image)
</code></pre>
<p>I tried inserting torch.cuda.synchronize() everywhere (before and after), but that didn’t help. Any advice on how to make it work?</p>
<p>p.s. I realize that a simple solution would be to pass i to the forward method of PCTL_layer, however it’s complicated because that layer is packaged into nn.Sequential list, so I don’t know how to do that either.</p>

ptrblck: <p>I’ve executed your code on a machine with 8 GPUs and got the following output:</p>
<pre><code class="lang-python"> model(torch.randn(8, 100, device='cuda'))

&gt;  gpu 0  calculate_running: True
 gpu 1  calculate_running: True
 gpu 2  calculate_running: True
 gpu 4  calculate_running: True
 gpu 3  calculate_running: True
 gpu 7  calculate_running: True
 gpu 6  calculate_running: True
 gpu 5  calculate_running: True
</code></pre>
<p>Which seems to indicates the attribute is properly set.</p>

michaelklachko: <p>Did you try more than one iteration? Because when I use 2 GPUs with 3 iterations I get the following:</p>
<pre><code class="lang-auto"> gpu 1  calculate_running: True
 gpu 0  calculate_running: True
 gpu 1  calculate_running: True
 gpu 0  calculate_running: True
 gpu 0  calculate_running: True
 gpu 1  calculate_running: True
</code></pre>
<p>But I want to see this:</p>
<pre><code class="lang-auto"> gpu 1  calculate_running: True
 gpu 0  calculate_running: True
</code></pre>

ptrblck: <p>Thanks for the follow-up.<br>
I’ve missed the change in this attribute and get the undesired behavior.<br>
Register the condition as a <code>BoolTensor</code> and it should work:</p>
<pre><code class="lang-python">self.calculate_running = torch.tensor(calculate_running, dtype=torch.bool)
</code></pre>

michaelklachko: <p>Does not work <img src="https://discuss.pytorch.org/images/emoji/apple/frowning.png?v=9" title=":frowning:" class="emoji" alt=":frowning:"></p>
<p>I tried both setting it to <code>False</code> and to <code>torch.tensor(False, dtype=torch.bool)</code>:</p>
<pre><code class="lang-auto">
class PCTL_Layer(nn.Module):
    def __init__(self, calculate_running=False):
        super(PCTL_Layer, self).__init__()
        self.register_buffer('running', torch.zeros(1))
        self.calculate_running = torch.tensor(calculate_running, dtype=torch.bool)

    def forward(self, input):
        if self.calculate_running:
            pctl, _ = torch.kthvalue(input.view(-1), int(input.numel() * 0.9))
            self.running = pctl
            print(' gpu {}  calculate_running: {}'.format(torch.cuda.current_device(), self.calculate_running))
            self.calculate_running = False
            #self.calculate_running = torch.tensor(False, dtype=torch.bool)

        return self.running
</code></pre>

michaelklachko: <p>Even when I synchronize GPUs like this:</p>
<pre><code class="lang-auto">    def forward(self, input):
        if self.calculate_running:
            pctl, _ = torch.kthvalue(input.view(-1), int(input.numel() * 0.9))
            self.running = pctl
            print(' gpu {}  calculate_running: {}'.format(torch.cuda.current_device(), self.calculate_running))
            torch.cuda.synchronize()
            self.calculate_running = False
            torch.cuda.synchronize()

        return self.running
</code></pre>
<p>It still does not work. Strange, isn’t it?</p>

ptrblck: <p>You are right. I missed the second output due to an error message, sorry.<br>
Thinking about it, it might make sense that this flag is not propagated using <code>DataParallel</code>.<br>
E.g. what would the expected result be, if the models on GPU0,1,2 set this flag to <code>False</code>, while all others keep it as <code>True</code>?</p>
<p>I think the cleanest approach would be to manipulate this flag after the forward pass manually using:</p>
<pre><code class="lang-python">_ = model(torch.randn(8, 100, device='cuda'))
model.module.percentile.calculate_running = False
</code></pre>
<p>which will make sure that each new copy of the model uses the new value.</p>

michaelklachko: <p>Great, boolean variable assignment works. Thank you. However, now <code>self.running</code> is not being assigned correctly. When I run the following code on 2 GPUs:</p>
<pre><code class="lang-auto">class PCTL_Layer(nn.Module):
    def __init__(self, calculate_running=False):
        super(PCTL_Layer, self).__init__()
        self.register_buffer('running', torch.zeros(1))
        self.calculate_running = calculate_running

    def forward(self, input):
        if self.calculate_running:
            pctl, _ = torch.kthvalue(input.view(-1), int(input.numel() * 0.9))
            self.running = pctl
            self.calculate_running = False
        print('calculate_running: {}   input: {}   running: {:.4f}'.format(self.calculate_running, input, self.running.item()))
        return self.running


class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.percentile = PCTL_Layer(calculate_running=True)

    def forward(self, x):
        x = self.percentile(x)
        return x


model = Model()
model = torch.nn.DataParallel(model).cuda()

for i in range(2):
    print('\nIteration', i)
    output = model(torch.randn(2, 4, device='cuda'))
    print('model.module.percentile.running: {}  model output: {}'.format(model.module.percentile.running, output))
    if i == 0:
        model.module.percentile.calculate_running = False
</code></pre>
<p>Here’s the output:</p>
<pre><code class="lang-auto">Iteration 0
calculate_running: False   input: tensor([[ 0.3747, -1.1507, -1.4812,  0.1900]], device='cuda:1')   running: 0.1900
calculate_running: False   input: tensor([[ 0.3961, -0.0903, -0.0330, -1.9596]], device='cuda:0')   running: -0.0330
/home/michael/miniconda2/envs/pt/lib/python3.7/site-packages/torch/nn/parallel/_functions.py:61: UserWarning: Was asked to gather along dimension 0, but all input tensors were scalars; will instead unsqueeze and return a vector.
  warnings.warn('Was asked to gather along dimension 0, but all '
model.module.percentile.running: tensor([0.], device='cuda:0')  model output: tensor([-0.0330,  0.1900], device='cuda:0')

Iteration 1
calculate_running: False   input: tensor([[ 0.6758, -0.2846,  0.9055, -0.0532]], device='cuda:0')   running: 0.0000
calculate_running: False   input: tensor([[-0.3517, -0.8317, -0.7304,  0.7370]], device='cuda:1')   running: 0.0000
model.module.percentile.running: tensor([0.], device='cuda:0')  model output: tensor([0., 0.], device='cuda:0')
</code></pre>
<p>Why <code>self.running</code> is being reset to 0 after the first iteration? In the worst case, I want each GPU to use its own value of running. But ideally I want it to be the mean across all GPUs.</p>
<p>I can fix it like this:</p>
<pre><code class="lang-auto">class PCTL_Layer(nn.Module):
    def __init__(self, calculate_running=False):
        super(PCTL_Layer, self).__init__()
        self.register_buffer('running', torch.zeros(1))
        self.calculate_running = calculate_running

    def forward(self, input):
        if self.calculate_running:
            pctl, _ = torch.kthvalue(input.view(-1), int(input.numel() * 0.9))
            self.running = pctl
            model.module.running_list.append(self.running)
            self.calculate_running = False

        print('calculate_running: {}   input: {}   running: {:.4f}'.format(self.calculate_running, input, self.running.item()))
        return self.running


class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.percentile = PCTL_Layer(calculate_running=True)
        self.running_list = []

    def forward(self, x):
        x = self.percentile(x)
        return x

model = Model()
model = torch.nn.DataParallel(model).cuda()

for i in range(3):
    print('\nIteration', i)
    output = model(torch.randn(2, 4, device='cuda'))
    print('model.module.percentile.running: {}'.format(model.module.percentile.running))
    if i == 0:
        model.module.percentile.calculate_running = False
        model.module.percentile.running = torch.tensor(model.module.running_list, device='cuda:0').mean()
    print('model output (running): {}  running_list: {}\n'.format(output, model.module.running_list))

</code></pre>
<p>Which produces the following output:</p>
<pre><code class="lang-auto">Iteration 0
calculate_running: False   input: tensor([[0.2759, 1.1083, 0.3523, 0.6275]], device='cuda:1')   running: 0.6275
calculate_running: False   input: tensor([[-1.5641,  0.3595,  1.7428, -0.5368]], device='cuda:0')   running: 0.3595
/home/michael/miniconda2/envs/pt/lib/python3.7/site-packages/torch/nn/parallel/_functions.py:61: UserWarning: Was asked to gather along dimension 0, but all input tensors were scalars; will instead unsqueeze and return a vector.
  warnings.warn('Was asked to gather along dimension 0, but all '
model.module.percentile.running: tensor([0.], device='cuda:0')
model output (running): tensor([0.3595, 0.6275], device='cuda:0')  running_list: [tensor(0.3595, device='cuda:0'), tensor(0.6275, device='cuda:1')]


Iteration 1
calculate_running: False   input: tensor([[ 0.4074,  1.4806, -0.5506,  0.3985]], device='cuda:0')   running: 0.4935
calculate_running: False   input: tensor([[ 0.2667,  0.4528,  0.1397, -0.5707]], device='cuda:1')   running: 0.4935
model.module.percentile.running: 0.4935113787651062
model output (running): tensor([0.4935, 0.4935], device='cuda:0')  running_list: [tensor(0.3595, device='cuda:0'), tensor(0.6275, device='cuda:1')]
</code></pre>
<p>But this is pretty ugly. Is there a better way?</p>

michaelklachko: <p>Also, note that when I’m averaging the running_list I have to move it to gpu-0. How do other GPUs access it? I mean, do they copy it from GPU-0 every time, or just once? If every time, how can I distribute it manually?</p>

michaelklachko: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> any insight about keeping the variable on all GPUs?</p>

ptrblck: <p><code>nn.DataParallel</code> will scatter the model to all devices and gather their outputs on the master device. <a href="https://medium.com/huggingface/training-larger-batches-practical-tips-on-1-gpu-multi-gpu-distributed-setups-ec88c3e51255" rel="nofollow noopener">This blog post</a> gives a good overview.</p>
<p>Since you are assigning a new value to <code>running</code>, I think the best way would be to have a look how data parallel works <a href="https://github.com/pytorch/pytorch/blob/9d37179061ae577942a955fd2934972b9cfa5fb4/torch/nn/parallel/data_parallel.py#L151" rel="nofollow noopener">internally</a> and adapt these methods to your use case.</p>