CoCalc provides the best real-time collaborative environment for Jupyter Notebooks, LaTeX documents, and SageMath, scalable from individual users to large groups and classes!
CoCalc provides the best real-time collaborative environment for Jupyter Notebooks, LaTeX documents, and SageMath, scalable from individual users to large groups and classes!
Path: blob/main/prototype_source/maskedtensor_sparsity.py
Views: 494
# -*- coding: utf-8 -*-12"""3(Prototype) MaskedTensor Sparsity4=================================5"""67######################################################################8# Before working on this tutorial, please make sure to review our9# `MaskedTensor Overview tutorial <https://pytorch.org/tutorials/prototype/maskedtensor_overview.html>`.10#11# Introduction12# ------------13#14# Sparsity has been an area of rapid growth and importance within PyTorch; if any sparsity terms are confusing below,15# please refer to the `sparsity tutorial <https://pytorch.org/docs/stable/sparse.html>`__ for additional details.16#17# Sparse storage formats have been proven to be powerful in a variety of ways. As a primer, the first use case18# most practitioners think about is when the majority of elements are equal to zero (a high degree of sparsity),19# but even in cases of lower sparsity, certain formats (e.g. BSR) can take advantage of substructures within a matrix.20#21# .. note::22#23# At the moment, MaskedTensor supports COO and CSR tensors with plans to support additional formats24# (such as BSR and CSC) in the future. If you have any requests for additional formats,25# please file a feature request `here <https://github.com/pytorch/pytorch/issues>`__!26#27# Principles28# ----------29#30# When creating a :class:`MaskedTensor` with sparse tensors, there are a few principles that must be observed:31#32# 1. ``data`` and ``mask`` must have the same storage format, whether that's :attr:`torch.strided`, :attr:`torch.sparse_coo`, or :attr:`torch.sparse_csr`33# 2. ``data`` and ``mask`` must have the same size, indicated by :func:`size()`34#35# .. _sparse-coo-tensors:36#37# Sparse COO tensors38# ------------------39#40# In accordance with Principle #1, a sparse COO MaskedTensor is created by passing in two sparse COO tensors,41# which can be initialized by any of its constructors, for example :func:`torch.sparse_coo_tensor`.42#43# As a recap of `sparse COO tensors <https://pytorch.org/docs/stable/sparse.html#sparse-coo-tensors>`__, the COO format44# stands for "coordinate format", where the specified elements are stored as tuples of their indices and the45# corresponding values. That is, the following are provided:46#47# * ``indices``: array of size ``(ndim, nse)`` and dtype ``torch.int64``48# * ``values``: array of size `(nse,)` with any integer or floating point dtype49#50# where ``ndim`` is the dimensionality of the tensor and ``nse`` is the number of specified elements.51#52# For both sparse COO and CSR tensors, you can construct a :class:`MaskedTensor` by doing either:53#54# 1. ``masked_tensor(sparse_tensor_data, sparse_tensor_mask)``55# 2. ``dense_masked_tensor.to_sparse_coo()`` or ``dense_masked_tensor.to_sparse_csr()``56#57# The second method is easier to illustrate so we've shown that below, but for more on the first and the nuances behind58# the approach, please read the :ref:`Sparse COO Appendix <sparse-coo-appendix>`.59#6061import torch62from torch.masked import masked_tensor63import warnings6465# Disable prototype warnings and such66warnings.filterwarnings(action='ignore', category=UserWarning)6768values = torch.tensor([[0, 0, 3], [4, 0, 5]])69mask = torch.tensor([[False, False, True], [False, False, True]])70mt = masked_tensor(values, mask)71sparse_coo_mt = mt.to_sparse_coo()7273print("mt:\n", mt)74print("mt (sparse coo):\n", sparse_coo_mt)75print("mt data (sparse coo):\n", sparse_coo_mt.get_data())7677######################################################################78# Sparse CSR tensors79# ------------------80#81# Similarly, :class:`MaskedTensor` also supports the82# `CSR (Compressed Sparse Row) <https://pytorch.org/docs/stable/sparse.html#sparse-csr-tensor>`__83# sparse tensor format. Instead of storing the tuples of the indices like sparse COO tensors, sparse CSR tensors84# aim to decrease the memory requirements by storing compressed row indices.85# In particular, a CSR sparse tensor consists of three 1-D tensors:86#87# * ``crow_indices``: array of compressed row indices with size ``(size[0] + 1,)``. This array indicates which row88# a given entry in values lives in. The last element is the number of specified elements,89# while `crow_indices[i+1] - crow_indices[i]` indicates the number of specified elements in row i.90# * ``col_indices``: array of size ``(nnz,)``. Indicates the column indices for each value.91# * ``values``: array of size ``(nnz,)``. Contains the values of the CSR tensor.92#93# Of note, both sparse COO and CSR tensors are in a `beta <https://pytorch.org/docs/stable/index.html>`__ state.94#95# By way of example:96#9798mt_sparse_csr = mt.to_sparse_csr()99100print("mt (sparse csr):\n", mt_sparse_csr)101print("mt data (sparse csr):\n", mt_sparse_csr.get_data())102103######################################################################104# Supported Operations105# --------------------106#107# Unary108# ^^^^^109# All `unary operators <https://pytorch.org/docs/master/masked.html#unary-operators>`__ are supported, e.g.:110#111112mt.sin()113114######################################################################115# Binary116# ^^^^^^117# `Binary operators <https://pytorch.org/docs/master/masked.html#unary-operators>`__ are also supported, but the118# input masks from the two masked tensors must match. For more information on why this decision was made, please119# find our `MaskedTensor: Advanced Semantics tutorial <https://pytorch.org/tutorials/prototype/maskedtensor_advanced_semantics.html>`__.120#121# Please find an example below:122#123124i = [[0, 1, 1],125[2, 0, 2]]126v1 = [3, 4, 5]127v2 = [20, 30, 40]128m = torch.tensor([True, False, True])129130s1 = torch.sparse_coo_tensor(i, v1, (2, 3))131s2 = torch.sparse_coo_tensor(i, v2, (2, 3))132mask = torch.sparse_coo_tensor(i, m, (2, 3))133134mt1 = masked_tensor(s1, mask)135mt2 = masked_tensor(s2, mask)136137print("mt1:\n", mt1)138print("mt2:\n", mt2)139140######################################################################141#142143print("torch.div(mt2, mt1):\n", torch.div(mt2, mt1))144print("torch.mul(mt1, mt2):\n", torch.mul(mt1, mt2))145146######################################################################147# Reductions148# ^^^^^^^^^^149# Finally, `reductions <https://pytorch.org/docs/master/masked.html#reductions>`__ are supported:150#151152mt153154######################################################################155#156157print("mt.sum():\n", mt.sum())158print("mt.sum(dim=1):\n", mt.sum(dim=1))159print("mt.amin():\n", mt.amin())160161######################################################################162# MaskedTensor Helper Methods163# ^^^^^^^^^^^^^^^^^^^^^^^^^^^164# For convenience, :class:`MaskedTensor` has a number of methods to help convert between the different layouts165# and identify the current layout:166#167# Setup:168#169170v = [[3, 0, 0],171[0, 4, 5]]172m = [[True, False, False],173[False, True, True]]174175mt = masked_tensor(torch.tensor(v), torch.tensor(m))176mt177178######################################################################179# :meth:`MaskedTensor.to_sparse_coo()` / :meth:`MaskedTensor.to_sparse_csr()` / :meth:`MaskedTensor.to_dense()`180# to help convert between the different layouts.181#182183mt_sparse_coo = mt.to_sparse_coo()184mt_sparse_csr = mt.to_sparse_csr()185mt_dense = mt_sparse_coo.to_dense()186187######################################################################188# :meth:`MaskedTensor.is_sparse` -- this will check if the :class:`MaskedTensor`'s layout189# matches any of the supported sparse layouts (currently COO and CSR).190#191192print("mt_dense.is_sparse: ", mt_dense.is_sparse)193print("mt_sparse_coo.is_sparse: ", mt_sparse_coo.is_sparse)194print("mt_sparse_csr.is_sparse: ", mt_sparse_csr.is_sparse)195196######################################################################197# :meth:`MaskedTensor.is_sparse_coo()`198#199200print("mt_dense.is_sparse_coo(): ", mt_dense.is_sparse_coo())201print("mt_sparse_coo.is_sparse_coo: ", mt_sparse_coo.is_sparse_coo())202print("mt_sparse_csr.is_sparse_coo: ", mt_sparse_csr.is_sparse_coo())203204######################################################################205# :meth:`MaskedTensor.is_sparse_csr()`206#207208print("mt_dense.is_sparse_csr(): ", mt_dense.is_sparse_csr())209print("mt_sparse_coo.is_sparse_csr: ", mt_sparse_coo.is_sparse_csr())210print("mt_sparse_csr.is_sparse_csr: ", mt_sparse_csr.is_sparse_csr())211212######################################################################213# Appendix214# --------215#216# .. _sparse-coo-appendix:217#218# Sparse COO Construction219# ^^^^^^^^^^^^^^^^^^^^^^^220#221# Recall in our :ref:`original example <sparse-coo-tensors>`, we created a :class:`MaskedTensor`222# and then converted it to a sparse COO MaskedTensor with :meth:`MaskedTensor.to_sparse_coo`.223#224# Alternatively, we can also construct a sparse COO MaskedTensor directly by passing in two sparse COO tensors:225#226227values = torch.tensor([[0, 0, 3], [4, 0, 5]]).to_sparse()228mask = torch.tensor([[False, False, True], [False, False, True]]).to_sparse()229mt = masked_tensor(values, mask)230231print("values:\n", values)232print("mask:\n", mask)233print("mt:\n", mt)234235######################################################################236# Instead of using :meth:`torch.Tensor.to_sparse`, we can also create the sparse COO tensors directly,237# which brings us to a warning:238#239# .. warning::240#241# When using a function like :meth:`MaskedTensor.to_sparse_coo` (analogous to :meth:`Tensor.to_sparse`),242# if the user does not specify the indices like in the above example,243# then the 0 values will be "unspecified" by default.244#245# Below, we explicitly specify the 0's:246#247248i = [[0, 1, 1],249[2, 0, 2]]250v = [3, 4, 5]251m = torch.tensor([True, False, True])252values = torch.sparse_coo_tensor(i, v, (2, 3))253mask = torch.sparse_coo_tensor(i, m, (2, 3))254mt2 = masked_tensor(values, mask)255256print("values:\n", values)257print("mask:\n", mask)258print("mt2:\n", mt2)259260######################################################################261# Note that ``mt`` and ``mt2`` look identical on the surface, and in the vast majority of operations, will yield the same262# result. But this brings us to a detail on the implementation:263#264# ``data`` and ``mask`` -- only for sparse MaskedTensors -- can have a different number of elements (:func:`nnz`)265# **at creation**, but the indices of ``mask`` must then be a subset of the indices of ``data``. In this case,266# ``data`` will assume the shape of ``mask`` by ``data = data.sparse_mask(mask)``; in other words, any of the elements267# in ``data`` that are not ``True`` in ``mask`` (that is, not specified) will be thrown away.268#269# Therefore, under the hood, the data looks slightly different; ``mt2`` has the "4" value masked out and ``mt``270# is completely without it. Their underlying data has different shapes,271# which would make operations like ``mt + mt2`` invalid.272#273274print("mt data:\n", mt.get_data())275print("mt2 data:\n", mt2.get_data())276277######################################################################278# .. _sparse-csr-appendix:279#280# Sparse CSR Construction281# ^^^^^^^^^^^^^^^^^^^^^^^282#283# We can also construct a sparse CSR MaskedTensor using sparse CSR tensors,284# and like the example above, this results in a similar treatment under the hood.285#286287crow_indices = torch.tensor([0, 2, 4])288col_indices = torch.tensor([0, 1, 0, 1])289values = torch.tensor([1, 2, 3, 4])290mask_values = torch.tensor([True, False, False, True])291292csr = torch.sparse_csr_tensor(crow_indices, col_indices, values, dtype=torch.double)293mask = torch.sparse_csr_tensor(crow_indices, col_indices, mask_values, dtype=torch.bool)294mt = masked_tensor(csr, mask)295296print("mt:\n", mt)297print("mt data:\n", mt.get_data())298299######################################################################300# Conclusion301# ----------302# In this tutorial, we have introduced how to use :class:`MaskedTensor` with sparse COO and CSR formats and303# discussed some of the subtleties under the hood in case users decide to access the underlying data structures304# directly. Sparse storage formats and masked semantics indeed have strong synergies, so much so that they are305# sometimes used as proxies for each other (as we will see in the next tutorial). In the future, we certainly plan306# to invest and continue developing in this direction.307#308# Further Reading309# ---------------310#311# To continue learning more, you can find our312# `Efficiently writing "sparse" semantics for Adagrad with MaskedTensor tutorial <https://pytorch.org/tutorials/prototype/maskedtensor_adagrad.html>`__313# to see an example of how MaskedTensor can simplify existing workflows with native masking semantics.314#315316317