When to densify sparse tensors?

mkarikom: <p>For a dataset that uses sparse tensors, should I be calling <code>tensor.to_dense()</code> in dataset.<strong>getitem</strong>() or in my collate fn?</p>
<h4>
<a name="background-1" class="anchor" href="#background-1"></a>Background:</h4>
<p>I’ve noticed that making datasets that use [still beta] sparse tensors leads to a lot of this:</p>
<pre><code class="lang-auto">NotImplementedError: Could not run 'aten::remainder.Tensor' with arguments from the 'SparseCUDA' backend. This could be because the operator doesn't exist for this backend, or was omitted during the selective/custom build process (if using custom build). If you are a Facebook employee using PyTorch on mobile, please visit https://fburl.com/ptmfixes for possible resolutions. 'aten::remainder.Tensor' is only available for these backends: [CPU, CUDA, HIP, MPS, IPU, XPU, HPU, VE, Meta, PrivateUse1, PrivateUse2, PrivateUse3, FPGA, ORT, Vulkan, Metal, QuantizedCPU, QuantizedCUDA, QuantizedHIP, QuantizedMPS, QuantizedIPU, QuantizedXPU, QuantizedHPU, QuantizedVE, QuantizedMeta, QuantizedPrivateUse1, QuantizedPrivateUse2, QuantizedPrivateUse3, CustomRNGKeyId, MkldnnCPU, SparseCsrCPU, SparseCsrCUDA, BackendSelect, Python, FuncTorchDynamicLayerBackMode, Functionalize, Named, Conjugate, Negative, ZeroTensor, ADInplaceOrView, AutogradOther, AutogradCPU, AutogradCUDA, AutogradHIP, AutogradXLA, AutogradMPS, AutogradIPU, AutogradXPU, AutogradHPU, AutogradVE, AutogradLazy, AutogradMeta, AutogradPrivateUse1, AutogradPrivateUse2, AutogradPrivateUse3, AutogradNestedTensor, Tracer, AutocastCPU, AutocastCUDA, FuncTorchBatched, FuncTorchVmapMode, Batched, VmapMode, FuncTorchGradWrapper, PythonTLSSnapshot, FuncTorchDynamicLayerFrontMode, PythonDispatcher].
</code></pre>

nivek: <p>I don’t think calling that in <code>__getitem__</code> vs collate function will get around the <code>NotImplementedError</code>.</p>
<p>When are you calling <code>remainder</code>? You might have to do <code>torch.remainder(sparse_coo1.to_dense(), sparse_coo2.to_dense())</code> but I can’t tell because I don’t know how you are using it.</p>

mkarikom: <p>Hi <a class="mention" href="/u/nivek">@nivek</a>, It turns out that I had a bug in my collate fn.  I have since fixed that but currently I am unable to set <code>pin_memory=True</code> in the DataLoader (I created a separate post specifically about this error <a href="https://discuss.pytorch.org/t/training-error-when-pin-memory-true-and-collate-fn-passes-sparse-tensors-to-the-batch/179713">here</a>).  I’d be very interested in any feedback you have on the training error described (at the very end) below!</p>
<h3>
<a name="dataset__getitem__-and-collate_fn-works-without-pin_memory-1" class="anchor" href="#dataset__getitem__-and-collate_fn-works-without-pin_memory-1"></a><code>DataSet:__getitem__()</code> and <code>collate_fn</code> (works without <code>pin_memory</code>)</h3>
<p>Here is <code>__getitem()</code>:</p>
<pre><code class="lang-auto">    def __getitem__(self,idx):
        if torch.is_tensor(idx):
            idx = idx.tolist()
        sample = {
            'rna':self.rna.__getitem__(idx),
            'atac':self.atac.__getitem__(idx),
            'cell_idx':idx,
            'batch_idx':self.mdata.obs.batch[idx],
            # 'cat_obs':pd.to_numeric(self.mdata.obs[self.cat_obs_names].iloc[idx])
            }
        for i in self.cat_obs_names:
            sample[i] = pd.to_numeric(self.mdata.obs[i].iloc[idx])
        return(sample)
</code></pre>
<p>Here is the <code>collate_fn</code> I pass to the DataLoader:</p>
<pre><code class="lang-auto">def sparse_batch_collate(batch:list,device:torch.device='cpu'):
    """
    Collate function to transform anndata csr view to pytorch sparse tensor
    """
    if type(batch[0]['atac'].X) == anndata._core.views.SparseCSRView:
        atac_batch = sparse_csr_to_tensor(vstack([x['atac'].X for x in batch])).to_dense().to(device)
    else:
        atac_batch = torch.FloatTensor(vstack([x['atac'].X for x in batch])).to(device)

    if type(batch[0]['rna'].X) == anndata._core.views.SparseCSRView:
        rna_batch = sparse_csr_to_tensor(vstack([x['rna'].X for x in batch])).to_dense().to(device)
    else:
        rna_batch = torch.FloatTensor(vstack([x['rna'].X for x in batch])).to(device)

    batch_idx = torch.tensor([[x['batch_idx']] for x in batch],dtype=torch.int16).to(device)
    cell_idx = torch.tensor([[x['cell_idx']] for x in batch],dtype=torch.int16).to(device)

    batch_collate = {
        'atac':atac_batch,
        'rna':rna_batch,
        'cell_idx':cell_idx,
        'batch_idx':batch_idx,
    }
    for key in batch[0].keys() - set(['rna','atac','cell_idx','batch_idx']):
        batch_collate[key] = torch.tensor([[x[key]] for x in batch],dtype=torch.int16).to(device)
    return(batch_collate)
</code></pre>
<h3>
<a name="working-pin_memoryfalse-2" class="anchor" href="#working-pin_memoryfalse-2"></a>Working: <code>pin_memory=False</code>
</h3>
<p>The following DataLoader works and my model trains:</p>
<pre><code class="lang-auto">loader = DataLoader(
    mds,
    batch_size=batch_size,
    collate_fn = lambda b: td.sparse_batch_collate(b, device=device),
)
</code></pre>
<h3>
<a name="not-working-pin_memorytrue-3" class="anchor" href="#not-working-pin_memorytrue-3"></a>Not Working: <code>pin_memory=True</code>
</h3>
<p>However, setting <code>pin_memory=True</code> in the above DataLoader generates the following error during training:</p>
<pre><code class="lang-auto">Epoch 1 of 10
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
Cell In[11], line 20
     18 for epoch in range(epochs):
     19     print(f"Epoch {epoch+1} of {epochs}")
---&gt; 20     train_epoch_loss = fit(model, loader)
     21     train_loss.append(train_epoch_loss)
     22     print(f"Train Loss: {train_epoch_loss:.4f}")

Cell In[11], line 4, in fit(model, dataloader)
      1 def fit(model, dataloader):
      2     # model.train()
      3     running_loss = 0.0
----&gt; 4     for batch in dataloader:
      5         optimizer.zero_grad()
      6         # reconstruction, mu, logvar = model(batch[0],batch[1])
      7         # reconstruction = model(batch[0],batch[1])

File /usr/local/python/lib/python3.8/site-packages/torch/utils/data/dataloader.py:628, in _BaseDataLoaderIter.__next__(self)
    625 if self._sampler_iter is None:
    626     # TODO(https://github.com/pytorch/pytorch/issues/76750)
    627     self._reset()  # type: ignore[call-arg]
--&gt; 628 data = self._next_data()
    629 self._num_yielded += 1
    630 if self._dataset_kind == _DatasetKind.Iterable and \
    631         self._IterableDataset_len_called is not None and \
    632         self._num_yielded &gt; self._IterableDataset_len_called:

File /usr/local/python/lib/python3.8/site-packages/torch/utils/data/dataloader.py:673, in _SingleProcessDataLoaderIter._next_data(self)
    671 data = self._dataset_fetcher.fetch(index)  # may raise StopIteration
    672 if self._pin_memory:
--&gt; 673     data = _utils.pin_memory.pin_memory(data, self._pin_memory_device)
    674 return data

File /usr/local/python/lib/python3.8/site-packages/torch/utils/data/_utils/pin_memory.py:58, in pin_memory(data, device)
     56 elif isinstance(data, collections.abc.Mapping):
     57     try:
---&gt; 58         return type(data)({k: pin_memory(sample, device) for k, sample in data.items()})  # type: ignore[call-arg]
     59     except TypeError:
     60         # The mapping type may not support `__init__(iterable)`.
     61         return {k: pin_memory(sample, device) for k, sample in data.items()}

File /usr/local/python/lib/python3.8/site-packages/torch/utils/data/_utils/pin_memory.py:58, in (.0)
     56 elif isinstance(data, collections.abc.Mapping):
     57     try:
---&gt; 58         return type(data)({k: pin_memory(sample, device) for k, sample in data.items()})  # type: ignore[call-arg]
     59     except TypeError:
     60         # The mapping type may not support `__init__(iterable)`.
     61         return {k: pin_memory(sample, device) for k, sample in data.items()}

File /usr/local/python/lib/python3.8/site-packages/torch/utils/data/_utils/pin_memory.py:53, in pin_memory(data, device)
     51 def pin_memory(data, device=None):
     52     if isinstance(data, torch.Tensor):
---&gt; 53         return data.pin_memory(device)
     54     elif isinstance(data, string_classes):
     55         return data

RuntimeError: cannot pin 'torch.cuda.FloatTensor' only dense CPU tensors can be pinned
</code></pre>

nivek: <p>It seems like your <code>collate_fn</code> may be moving data to <code>device=cuda</code>. Once you move it there, you can’t pin it anymore, as the error indicates:<br>
<code>RuntimeError: cannot pin 'torch.cuda.FloatTensor' only dense CPU tensors can be pinned</code></p>
<p>You can only pin objects that are on CPUs.</p>

mkarikom: <p>Thanks <a class="mention" href="/u/nivek">@nivek</a>, can you comment on the relative efficiency of the following two possible strategies?:</p>
<ol>
<li>densify the sparse CPU tensors during collate (pin_memory=True) → move them to the GPU during training loop</li>
<li>move the sparse CPU tensors to the GPU during collate (pin_memory=False) → densify as needed</li>
</ol>

nivek: <p>Unfortunately I do not know the exact performance numbers, but I do know that it is not recommended to move data to GPU with multiprocessing because of the many subtleties involved. You can read more <a href="https://pytorch.org/docs/stable/notes/multiprocessing.html#multiprocessing-cuda-note" rel="noopener nofollow ugc">about it here</a>.</p>