介紹下pytorch的主要模塊和輔助模塊。
Pytorch的主要模塊
PyTorch主要包括以下16個模塊：
1．torch模塊
torch模塊本身包含了PyTorch經常使用的一些激活函數，比如Sigmoid（torch.sigmoid）、ReLU（torch.relu）和Tanh（torch.tanh），以及PyTorch張量的一些操作，比如矩陣的乘法（torch.mm）、張量元素的選擇（torch.select）。需要注意的是，這些操作的對象大多數都是張量，因此，傳入的參數需要是PyTorch的張量，否則會報錯（一般報類型錯誤，即TypeError）。另外，還有一類函數能夠產生一定形狀的張量，比如torch.zeros產生元素全為0的張量，torch.randn產生元素服從標准正態分布的張量等。
2．torch.Tensor模塊
torch.Tensor模塊定義了torch中的張量類型，其中的張量有不同的數值類型，如單精度、雙精度浮點、整數類型等，而且張量有一定的維數和形狀。同時，張量的類中也包含着一系列的方法，返回新的張量或者更改當前的張量。torch.Storage則負責torch.Tensor底層的數據存儲，即前面提到的為一個張量分配連續的一維內存地址（用於存儲相同類型的一系列元素，數目則為張量的總元素數目）。這里需要提到一點，如果張量的某個類方法會返回張量，按照PyTorch中的命名規則，如果張量方法后綴帶下畫線，則該方法會修改張量本身的數據，反之則會返回新的張量。比如，Tensor.add方法會讓當前張量和輸入參數張量做加法，返回新的張量，而Tensor.add_方法會改變當前張量的值，新的值為舊的值和輸入參數之和。
3．torch.sparse模塊
torch.sparse模塊定義了稀疏張量，其中構造的稀疏張量采用的是COO格式（Coordinate），主要方法是用一個長整型定義非零元素的位置，用浮點數張量定義對應非零元素的值。稀疏張量之間可以做元素加、減、乘、除運算和矩陣乘法。
4．torch.cuda模塊
torch.cuda模塊定義了與CUDA運算相關的一系列函數，包括但不限於檢查系統的CUDA是否可用，當前進程對應的GPU序號（在多GPU情況下），清除GPU上的緩存，設置GPU的計算流（Stream），同步GPU上執行的所有核函數（Kernel）等。
5．torch.nn模塊
torch.nn是一個非常重要的模塊，是PyTorch神經網絡模塊化的核心。這個模塊定義了一系列模塊，包括卷積層nn.ConvNd（N=1，2，3）和線性層（全連接層）nn.Linear等。當構建深度學習模型的時候，可以通過繼承nn.Module類並重寫forward方法來實現一個新的神經網絡（后續會提到如何通過組合神經網絡模塊來構建深度學習模型）。另外，torch.nn中也定義了一系列的損失函數，包括平方損失函數（torch.nn.MSELoss）、交叉熵損失函數（torch.nn.CrossEntropyLoss）等。一般來說，torch.nn里定義的神經網絡模塊都含有參數，可以對這些參數使用優化器進行訓練。
6．torch.nn.functional函數模塊
torch.nn.functional是PyTorch的函數模塊，定義了一些核神經網絡相關的函數，包括卷積函數和池化函數等，這些函數也是深度學習模型構建的基礎。需要指出的是，torch.nn中定義的模塊一般會調用torch.nn.functional里的函數，比如，nn.ConvNd模塊（N=1，2，3）會調用torch.nn.functional.convNd函數（N=1，2，3）。另外，torch.nn.functional里面還定義了一些不常用的激活函數，包括torch.nn.functional.relu6和torch.nn.functional.elu等。
7．torch.nn.init模塊
torch.nn.init模塊定義了神經網絡權重的初始化。前面已經介紹過，如果初始的神經網絡權重取值不合適，就會導致后續的優化過程收斂很慢，甚至不熟練。這個模塊中的函數就是為了解決神經網絡權重的初始化問題，其中使用了很多初始化方法，包括均勻初始化torch.nn.init.uniform_和正態分布歸一化torch.nn.init.normal_等。在前面提到過，在PyTorch中函數或者方法如果以下畫線結尾，則這個方法會直接改變作用張量的值。因此，這些方法會直接改變傳入張量的值，同時會返回改變后的張量。
8．torch.optim模塊
torch.optim模塊定義了一系列的優化器，包括但不限於前一章介紹的優化器，如torch.optim.SGD（隨機梯度下降算法）、torch.optim.Adagrad（AdaGrad算法）、torch.optim.RMSprop（RMSProp算法）和torch.optim.Adam（Adam算法）等。當然，這個模塊還包含學習率衰減的算法的子模塊，即torch.optim.lr_scheduler，這個子模塊中包含了諸如學習率階梯下降算法torch.optim.lr_scheduler.StepLR和余弦退火算法torch.optim.lr_scheduler.CosineAnnealingLR等學習率衰減算法。
9．torch.autograd模塊
torch.autograd模塊是PyTorch的自動微分算法模塊，定義了一系列的自動微分函數，包括torch.autograd.backward函數，主要用於在求得損失函數之后進行反向梯度傳播，torch.autograd.grad函數用於一個標量張量（即只有一個分量的張量）對另一個張量求導，以及在代碼中設置不參與求導的部分。另外，這個模塊還內置了數值梯度功能和檢查自動微分引擎是否輸出正確結果的功能。
10．torch.distributed模塊
torch.distributed是PyTorch的分布式計算模塊，主要功能是提供PyTorch並行運行環境，其主要支持的后端有MPI、Gloo和NCCL三種。PyTorch的分布式工作原理主要是啟動多個並行的進程，每個進程都擁有一個模型的備份，然后輸入不同的訓練數據到多個並行的進程，計算損失函數，每個進程獨立地做反向傳播，最后對所有進程權重張量的梯度做歸約（Reduce）。用到后端的部分主要是數據的廣播（Broadcast）和數據的收集（Gather），其中，前者是把數據從一個節點（進程）傳播到另一個節點（進程），比如歸約后梯度張量的傳播，后者則是把數據從其他節點（進程）轉移到當前節點（進程），比如把梯度張量從其他節點轉移到某個特定的節點，然后對所有的張量求平均。PyTorch的分布式計算模塊不但提供了后端的一個包裝，還提供了一些啟動方式來啟動多個進程，包括但不限於通過網絡（TCP）、通過環境變量、通過共享文件等。
11．torch.distributions模塊
torch.distributions模塊提供了一系列類，使得PyTorch能夠對不同的分布進行采樣，並且生成概率采樣過程的計算圖。在一些應用過程中，比如強化學習（Reinforcement Learning），經常會使用一個深度學習模型來模擬在不同環境條件下采取的策略（Policy），其最后的輸出是不同動作的概率。當深度學習模型輸出概率之后，需要根據概率對策略進行采樣來模擬當前的策略概率分布，最后用梯度下降方法來讓最優策略的概率最大（這個算法稱為策略梯度算法，PolicyGradient）。實際上，因為采樣的輸出結果是離散的，無法直接求導，所以不能使用反向傳播的方法來優化網絡。torch.distributions模塊的存在目的就是為了解決這個問題。我們可以結合torch.distributions.Categorical進行采樣，然后使用對數求導技巧來規避這個問題。當然，除了服從多項式分布的torch.distributions.Categorical類，PyTorch還支持其他的分布（包括連續分布和離散分布），比如torch.distributions.Normal類支持連續的正態分布的采樣，可以用於連續的強化學習的策略。
12．torch.hub模塊
torch.hub提供了一系列預訓練的模型供用戶使用。比如，可以通過torch.hub.list函數來獲取某個模型鏡像站點的模型信息。通過torch.hub.load來載入預訓練的模型，載入后的模型可以保存到本地，並可以看到這些模型對應類支持的方法。更多torch.hub支持的模型可以參考PyTorch官網中的相關頁面。
13．torch.jit模塊
torch.jit是PyTorch的即時編譯器（Just-In-Time Compiler，JIT）模塊。這個模塊存在的意義是把PyTorch的動態圖轉換成可以優化和序列化的靜態圖，其主要工作原理是通過輸入預先定義好的張量，追蹤整個動態圖的構建過程，得到最終構建出來的動態圖，然后轉換為靜態圖（通過中間表示，即IntermediateRepresentation，來描述最后得到的圖）。通過JIT得到的靜態圖可以被保存，並且被PyTorch其他的前端（如C++語言的前端）支持。另外，JIT也可以用來生成其他格式的神經網絡描述文件，如前文敘述的ONNX。需要注意的一點是，torch.jit支持兩種模式，即腳本模式（ScriptModule）和追蹤模式（Tracing）。前者和后者都能構建靜態圖，區別在於前者支持控制流，后者不支持，但是前者支持的神經網絡模塊比后者少，比如腳本模式不支持torch.nn.GRU（詳細的描述可以參考PyTorch官方提供的JIT相關的文檔）。
14．torch.multiprocessing模塊
torch.multiprocessing定義了PyTorch中的多進程API。通過使用這個模塊，可以啟動不同的進程，每個進程運行不同的深度學習模型，並且能夠在進程間共享張量（通過共享內存的方式）。共享的張量可以在CPU上，也可以在GPU上，多進程API還提供了與Python原生的多進程API（即multiprocessing庫）相同的一系列函數，包括鎖（Lock）和隊列（Queue）等。
15．torch.random模塊
torch.random提供了一系列的方法來保存和設置隨機數生成器的狀態，包括使用get_rng_state函數獲取當前隨機數生成器狀態，set_rng_state函數設置當前隨機數生成器狀態，並且可以使用manual_seed函數來設置隨機種子，也可使用initial_seed函數來得到程序初始的隨機種子。因為神經網絡的訓練是一個隨機的過程，包括數據的輸入、權重的初始化都具有一定的隨機性。設置一個統一的隨機種子可以有效地幫助我們測試不同結構神經網絡的表現，有助於調試神經網絡的結構
16．torch.onnx模塊
torch.onnx定義了PyTorch導出和載入ONNX格式的深度學習模型描述文件。前面已經介紹過，ONNX格式的存在是為了方便不同深度學習框架之間交換模型。引入這個模塊可以方便PyTorch導出模型給其他深度學習框架使用，或者讓PyTorch可以載入其他深度學習框架構建的深度學習模型。
Pytorch的輔助模塊
1．torch.utils.bottleneck模塊
torch.utils.bottleneck可以用來檢查深度學習模型中模塊的運行時間，從而可以找到導致性能瓶頸的那些模塊，通過優化那些模塊的運行時間，從而優化整個深度學習模型的性能。
2．torch.utils.checkpoint模塊
torch.utils.checkpoint可以用來節約深度學習使用的內存。通過前面的介紹我們知道，因為要進行梯度反向傳播，在構建計算圖的時候需要保存中間的數據，而這些數據大大增加了深度學習的內存消耗。為了減少內存消耗，讓迷你批次的大小得到提高，從而提升深度學習模型的性能和優化時的穩定性，我們可以通過這個模塊記錄中間數據的計算過程，然后丟棄這些中間數據，等需要用到的時候再重新計算這些數據。這個模塊設計的核心思想是以計算時間換內存空間，當使用得當的時候，深度學習模型的性能可以有很大的提升。
3．torch.utils.cpp_extension模塊
torch.utils.cpp_extension定義了PyTorch的C++擴展，其主要包含兩個類：CppExtension定義了使用C++來編寫的擴展模塊的源代碼相關信息，CUDAExtension則定義了C++/CUDA編寫的擴展模塊的源代碼相關信息。在某些情況下，用戶可能需要使用C++實現某些張量運算和神經網絡結構（比如PyTorch沒有類似功能的模塊或者PyTorch類似功能的模塊性能比較低），PyTorch的C++擴展模塊就提供了一個方法能夠讓Python來調用使用C++/CUDA編寫的深度學習擴展模塊。在底層上，這個擴展模塊使用了pybind11，保持了接口的輕量性並使得PyTorch易於被擴展。在后續章節會介紹如何使用C++/CUDA來編寫PyTorch的擴展。
4．torch.utils.data模塊
torch.utils.data引入了數據集（Dataset）和數據載入器（DataLoader）的概念，前者代表包含了所有數據的數據集，通過索引能夠得到某一條特定的數據，后者通過對數據集的包裝，可以對數據集進行隨機排列（Shuffle）和采樣（Sample），得到一系列打亂數據順序的迷你批次。
5．torch.utils.dlpacl模塊
torch.utils.dlpack定義了PyTorch張量和DLPack張量存儲格式之間的轉換，用於不同框架之間張量數據的交換。
6．torch.utils.tensorboard模塊
torch.utils.tensorboard是PyTorch對TensorBoard數據可視化工具的支持。TensorBoard原來是TensorFlow自帶的數據可視化工具，能夠顯示深度學習模型在訓練過程中損失函數、張量權重的直方圖，以及模型訓練過程中輸出的文本、圖像和視頻等。TensorBoard的功能非常強大，而且是基於可交互的動態網頁設計的，使用者可以通過預先提供的一系列功能來輸出特定的訓練過程的細節（如某一神經網絡層的權重的直方圖，以及訓練過程中某一段時間的損失函數等）。PyTorch支持TensorBoard可視化之后，在PyTorch的訓練過程中，可以很方便地觀察中間輸出的張量，也可以方便地調試深度學習模型。