Using Pytorch model for data augment inside of __getitem__

Archjbald: <p>Hello,</p>
<p>I’m trying to implement a data augmentation method using a deep model. To do so, I initialized the augmentation network in the <strong>init</strong> of my Dataset class with trained weights and stored it as an attribute.</p>
<p>However, when the code enters the <strong>getitem</strong> method of the dataset, all the weights and stored Tensor of my augmentation network are reset to 0, as if it was never initialized.</p>
<p>Am I doing something wrong? I want to use the network as the first step of the data augmentation process, so I can’t put it outside of the <strong>getitem</strong> method.</p>
<p>Thank you for your insights!</p>

ptrblck: <p>Your idea seems to work for me using this simple approach:</p>
<pre><code class="lang-python">class MyDataset(Dataset):
    def __init__(self):
        self.model = nn.Linear(10, 10)
        print("Initial weight.abs().sum (): ", self.model.weight.abs().sum())
        self.data = torch.randn(10, 10)
        
    def __len__(self):
        return len(self.data)
    
    def __getitem__(self, index):
        print("Inside __getitem__ weight.abs().sum (): ", self.model.weight.abs().sum())
        x = self.data[index]
        x = x.unsqueeze(0)
        with torch.no_grad():
            x = self.model(x)
        return x
        
dataset = MyDataset()
# Initial weight.abs().sum ():  tensor(16.0544, grad_fn=&lt;SumBackward0&gt;)

for data in dataset:
    pass
# Inside __getitem__ weight.abs().sum ():  tensor(16.0544, grad_fn=&lt;SumBackward0&gt;)
# Inside __getitem__ weight.abs().sum ():  tensor(16.0544, grad_fn=&lt;SumBackward0&gt;)
# Inside __getitem__ weight.abs().sum ():  tensor(16.0544, grad_fn=&lt;SumBackward0&gt;)
# Inside __getitem__ weight.abs().sum ():  tensor(16.0544, grad_fn=&lt;SumBackward0&gt;)
# Inside __getitem__ weight.abs().sum ():  tensor(16.0544, grad_fn=&lt;SumBackward0&gt;)
# Inside __getitem__ weight.abs().sum ():  tensor(16.0544, grad_fn=&lt;SumBackward0&gt;)
# Inside __getitem__ weight.abs().sum ():  tensor(16.0544, grad_fn=&lt;SumBackward0&gt;)
# Inside __getitem__ weight.abs().sum ():  tensor(16.0544, grad_fn=&lt;SumBackward0&gt;)
# Inside __getitem__ weight.abs().sum ():  tensor(16.0544, grad_fn=&lt;SumBackward0&gt;)
# Inside __getitem__ weight.abs().sum ():  tensor(16.0544, grad_fn=&lt;SumBackward0&gt;)
# Inside __getitem__ weight.abs().sum ():  tensor(16.0544, grad_fn=&lt;SumBackward0&gt;)


loader = DataLoader(dataset, num_workers=2, batch_size=5)
for data in loader:
    pass
# output is interleaved due to num_workers=2
# Inside __getitem__ weight.abs().sum (): Inside __getitem__ weight.abs().sum ():   tensor(16.0544, grad_fn=&lt;SumBackward0&gt;)tensor(16.0544, grad_fn=&lt;SumBackward0&gt;)
# Inside __getitem__ weight.abs().sum (): Inside __getitem__ weight.abs().sum ():   tensor(16.0544, grad_fn=&lt;SumBackward0&gt;)tensor(16.0544, grad_fn=&lt;SumBackward0&gt;)
# Inside __getitem__ weight.abs().sum (): Inside __getitem__ weight.abs().sum ():   tensor(16.0544, grad_fn=&lt;SumBackward0&gt;)tensor(16.0544, grad_fn=&lt;SumBackward0&gt;)
# Inside __getitem__ weight.abs().sum (): Inside __getitem__ weight.abs().sum ():   tensor(16.0544, grad_fn=&lt;SumBackward0&gt;)tensor(16.0544, grad_fn=&lt;SumBackward0&gt;)
# Inside __getitem__ weight.abs().sum ():  Inside __getitem__ weight.abs().sum (): tensor(16.0544, grad_fn=&lt;SumBackward0&gt;) 
# tensor(16.0544, grad_fn=&lt;SumBackward0&gt;)
</code></pre>

Archjbald: <p>Thank you for your reply!</p>
<p>After further investigations (I thought I had already tried that but apparently not), my issue seems related to the DataLoader workers. When set to 0, it works perfectly but doesn’t with several workers.</p>
<p>Is there a particular reason why the parallel versions would not properly copy the model’s weights?</p>

ptrblck: <p>Multiple workers are fine in my code and I don’t know why you are seeing the issue.<br>
Is my code also breaking in your environment?</p>

Archjbald: <p>I might have figured out the issue: it is caused by the fact that my augmentation network was on GPU. Your code works perfectly fine until I transfer the Linear layer to the GPU.<br>
Does this mean I have to perform either on a CPU or with no parallel data worker? I don’t know if it is an expected behavior, I don’t know much about parallel computing.</p>
<p>Thank you again for your help.</p>