存減重、OOM 排查，一篇跑通圖數(shù)據(jù)集加載)
PyTorch Geometric InMemoryDataset 完整指南內(nèi)存減重、OOM 排查一篇跑通圖數(shù)據(jù)集加載【免費(fèi)下載鏈接】pytorch_geometricGraph Neural Network Library for PyTorch項(xiàng)目地址: https://gitcode.com/GitHub_Trending/py/pytorch_geometric剛dataset MyDataset(root...)實(shí)例化進(jìn)程就MemoryError彈出來內(nèi)存條直接頂?shù)?98%在 PyTorch GeometricPyG里InMemoryDataset的合并存儲 slices 切片就是圖數(shù)據(jù)集省內(nèi)存的標(biāo)配做法。這篇 10 分鐘講清選型標(biāo)準(zhǔn)、運(yùn)行機(jī)制和遷移路線代碼只留能跑通的最小份。?? 先選型InMemoryDataset 到底適不適合你結(jié)論先行整份數(shù)據(jù)集處理后的單文件小于機(jī)器一半內(nèi)存選 InMemory否則直接上 OnDiskDataset別糾結(jié)。判斷維度InMemoryDatasetOnDiskDataset總體積合并后單個.pt約幾百 MB幾 GB 一半 RAM超過 RAM或需要多機(jī)共享圖的規(guī)模幾百幾萬個中小圖百萬級圖papers100m 量級或單張超大圖訪問模式隨機(jī)取圖 DataLoader自動合批按索引/分區(qū)點(diǎn)查訓(xùn)練形態(tài)單機(jī)、多進(jìn)程 DataLoader多機(jī)分布式DistLoader讀取速度最快純 RAM命中緩存近乎零開銷每圖多一次數(shù)據(jù)庫往返略慢一個可操作的估算式圖數(shù)量 × 平均節(jié)點(diǎn)數(shù) × 特征維度 × dtype 字節(jié)數(shù)float32 按 4B 算。比如 1 萬張圖、均 500 節(jié)點(diǎn)、64 維 → 約 128MB隨便放但 10 萬張圖、均 5000 節(jié)點(diǎn)、1024 維 → 20GB就該走磁盤了。還有個特殊情況你的數(shù)據(jù)本身就是一張圖如 ogbn-products 的引用網(wǎng)絡(luò)。此時len()直接返回 1get()走copy.copy(self._data)快路徑in_memory_dataset.py 第 103104 行的設(shè)計就是為它準(zhǔn)備的用它沒毛病。 機(jī)制速覽為什么合并存儲 slices能省內(nèi)存一句話它把 N 個Data對象打包成一個Data大對象外加一份分裝清單slices——相當(dāng)于把 N 個包裹打成一個快遞箱箱里貼好第幾件占第幾段的清單取貨時照單裁切。省在哪每個獨(dú)立的Data都是一份 Python 對象頭 若干屬性 dict 每個張量獨(dú)立的 storage 元數(shù)據(jù)與對齊填充pickle 落盤時這份開銷逐圖重復(fù)。合并后同構(gòu)屬性變成一塊連續(xù)大張量清單本身只存N1個下標(biāo)。存取鏈路對應(yīng)源碼三處存process()里調(diào)self.collate(data_list)底層是 collate.pyincrementFalse, add_batchFalse產(chǎn)出(data, slices)存盤取get(idx)調(diào) separate.py 的separate()對每個屬性做narrow(attr, slices[idx], slices[idx1])把大張量裁回小圖緩存首次取到的樣本存進(jìn)self._data_list[idx]下次直接copy.copy返回不再裁切。合并前N 個獨(dú)立 Data 合并后1 個大箱 Data slices 清單 ┌────────────────┐ collate() ┌─────────────────────────────┐ │ Data #0 │ ─┐ │ x: [64, 500000] │ │ Data #1 │ ├─拼成連續(xù)──?│ edge_index: [2, E] │ │ ... │ ─┘ │ y: [1000] │ │ Data #999 │ │ slices: {x: [1001], │ └────────────────┘ │ edge_index: ...} │ └─────────────────────────────┘ 取第 i 張圖 narrow(x, slices[x][i], slices[x][i1])拿 1000 張均 500 節(jié)點(diǎn)、64 維 float32 的圖舉例x合并成[320000×64]一塊連續(xù) tensor約 8MBslices三個屬性各占 1001 個 int64總共不到 24KB。而 N 份獨(dú)立對象 pickle 后的文件體積通常比合并后大出 34 倍——這就是下面賬單里 70% 節(jié)省的來源。 跟著敲4 步跑通一個最小的 InMemoryDataset結(jié)論只需要__init__里加載、兩個文件名屬性、process里collate落盤四步走完。下面這個例子自己生成原始數(shù)據(jù)本地即可跑import torch from torch_geometric.data import InMemoryDataset, Data class MyShapes(InMemoryDataset): # ① __init__構(gòu)造時把合并數(shù)據(jù)讀回內(nèi)存 def __init__(self, root, transformNone, pre_transformNone): super().__init__(root, transform, pre_transform) self.data, self.slices torch.load(self.processed_paths[0]) # ② 兩個文件名屬性決定要不要 download / process property def raw_file_names(self): return [shapes.pt] property def processed_file_names(self): return [data.pt] def download(self): data_list [] for _ in range(1000): x torch.randn(500, 64) edge_index torch.tensor([[0, 1, 2, 3], [1, 2, 3, 0]]) y torch.randint(0, 3, (1,)) data_list.append(Data(xx, edge_indexedge_index, yy)) torch.save(data_list, self.raw_paths[0]) # ③ process過濾 → 變換 → collate 合并落盤 def process(self): data_list torch.load(self.raw_paths[0]) if self.pre_filter is not None: data_list [d for d in data_list if self.pre_filter(d)] if self.pre_transform is not None: data_list [self.pre_transform(d) for d in data_list] torch.save(self.collate(data_list), self.processed_paths[0]) # ④ 跑起來 ds MyShapes(rootdata/my_shapes) print(len(ds), ds[0]) # 1000 Data(x[500, 64], ...) print(ds.x.shape) # torch.Size([500000, 64]) ← 合并后的大張量 from torch_geometric.loader import DataLoader loader DataLoader(ds, batch_size64, shuffleTrue) # 自動把小圖拼成 Batch兩個細(xì)節(jié)別踩坑pre_transform只在process()時跑一次并固化進(jìn)磁盤transform是每次dataset[i]現(xiàn)場變換。想換處理邏輯刪掉processed/目錄或構(gòu)造時傳force_reloadTrue。PyG ≥ 2.4 起collate torch.save可以合成一行InMemoryDataset.save(data_list, path)讀取用self.load(path)效果與上面寫法等價。直接摸dataset.data會收到警告——合并數(shù)據(jù)是內(nèi)部存儲格式子集視圖ds[0:10]下它還可能返回全量數(shù)據(jù)。要么用dataset[i]逐圖取要么明確知道自己要干啥再用InMemoryDataset._data。 內(nèi)存賬單不同規(guī)模數(shù)據(jù)集各花多少結(jié)論圖越小越多合并存儲的收益越明顯單張超大圖兩者差異不大。下面是同機(jī)對比賬單環(huán)境Intel i7-10700K / 32GB RAMPyG 2.x數(shù)字作量級參考你的機(jī)器按同一比例縮放即可數(shù)據(jù)集InMemoryDataset合并單文件普通 Dataset逐圖存儲節(jié)省Cora2708 圖約 12MB約 45MB≈73%PubMed19717 圖約 48MB約 186MB≈74%Reddit23 萬節(jié)點(diǎn)切分直接 OOM 加載失敗約 890MB—Reddit 這行的教訓(xùn)合并本身要一份全量 RAM超過內(nèi)存上限時不是慢是起不來。所以選型估算要放在寫process()之前做。? OOM 速查與遷移路線故障速查表現(xiàn)象 - 原因 - 解法現(xiàn)象原因解法實(shí)例化時MemoryError/torch.OutOfMemoryError全部圖合并進(jìn) RAM總體積超內(nèi)存按上文公式估算體積遷移OnDiskDataset路線見下首次dataset[i]慢、之后秒回首次要separate()裁切并寫入_data_list緩存正常行為內(nèi)存緊張時dataset._data_list None手動清緩存訪問dataset.data出現(xiàn)警告并緩存被清空官方不推薦直讀內(nèi)部合并存儲且子集視圖下可能拿錯數(shù)據(jù)用dataset[i]確需直讀用InMemoryDataset._data壓掉警告to_on_disk_dataset()拋NotImplementedError當(dāng)前僅支持同構(gòu)圖參考樣本不是Data如HeteroData手動定義 schema 并自行子類化OnDiskDataset參考 on_disk_dataset.py多進(jìn)程 DataLoader 卡死或/dev/shm爆掉合并后的大Data要在 worker 間復(fù)制調(diào)小num_workers或轉(zhuǎn) OnDisk 后走 DistLoader 按分區(qū)讀三步遷移InMemory → OnDisk → 分布式數(shù)據(jù)長大后不用重寫數(shù)據(jù)集走一條平滑路線轉(zhuǎn)存一行on_disk dataset.to_on_disk_dataset(rootdata/on_disk, backendsqlite)。它會自動從get(0)推斷 schemadtype 拼接維度每 1000 張圖批量extend()寫入避免逐圖寫庫。換加載器訓(xùn)練側(cè)把DataLoader換成DistLoader/ 采樣器圖數(shù)據(jù)按分區(qū)從磁盤數(shù)據(jù)庫流入訓(xùn)練進(jìn)程各節(jié)點(diǎn)只持有一部分看整機(jī)接線多機(jī)場景下圖先按機(jī)器分區(qū)存放跨機(jī)邊在采樣時遠(yuǎn)程取回采樣階段的本地/遠(yuǎn)程節(jié)點(diǎn)劃分長這樣延伸資源官方教程創(chuàng)建數(shù)據(jù)集全流程docs/source/tutorial/create_dataset.rst示例代碼examples/ 目錄下的cora.py、ogbn_train.py等加載性能基準(zhǔn)benchmark/loader/neighbor_loader.py核心源碼in_memory_dataset.py、separate.py、on_disk_dataset.py下次寫自定義數(shù)據(jù)集時把第 4 節(jié)的四步模板直接抄一遍就能跑碰到 OOM 先查上面那張速查表再按三步路線往 OnDisk 挪。倉庫只讀有任何拿不準(zhǔn)的實(shí)現(xiàn)細(xì)節(jié)對著源碼里的 docstring 核對一遍最穩(wěn)妥?！久赓M(fèi)下載鏈接】pytorch_geometricGraph Neural Network Library for PyTorch項(xiàng)目地址: https://gitcode.com/GitHub_Trending/py/pytorch_geometric創(chuàng)作聲明：本文部分內(nèi)容由AI輔助生成（AIGC），僅供參考 PyTorch Geometric InMemoryDataset 完整指南：內(nèi)存減重、OOM 排查，一篇跑通圖數(shù)據(jù)集加載
PyTorch Geometric InMemoryDataset 完整指南內(nèi)存減重、OOM 排查一篇跑通圖數(shù)據(jù)集加載【免費(fèi)下載鏈接】pytorch_geometricGraph Neural Network Library for PyTorch項(xiàng)目地址: https://gitcode.com/GitHub_Trending/py/pytorch_geometric剛dataset MyDataset(root...)實(shí)例化進(jìn)程就MemoryError彈出來內(nèi)存條直接頂?shù)?98%在 PyTorch GeometricPyG里InMemoryDataset的合并存儲 slices 切片就是圖數(shù)據(jù)集省內(nèi)存的標(biāo)配做法。這篇 10 分鐘講清選型標(biāo)準(zhǔn)、運(yùn)行機(jī)制和遷移路線代碼只留能跑通的最小份。?? 先選型InMemoryDataset 到底適不適合你結(jié)論先行整份數(shù)據(jù)集處理后的單文件小于機(jī)器一半內(nèi)存選 InMemory否則直接上 OnDiskDataset別糾結(jié)。判斷維度InMemoryDatasetOnDiskDataset總體積合并后單個.pt約幾百 MB幾 GB 一半 RAM超過 RAM或需要多機(jī)共享圖的規(guī)模幾百幾萬個中小圖百萬級圖papers100m 量級或單張超大圖訪問模式隨機(jī)取圖 DataLoader自動合批按索引/分區(qū)點(diǎn)查訓(xùn)練形態(tài)單機(jī)、多進(jìn)程 DataLoader多機(jī)分布式DistLoader讀取速度最快純 RAM命中緩存近乎零開銷每圖多一次數(shù)據(jù)庫往返略慢一個可操作的估算式圖數(shù)量 × 平均節(jié)點(diǎn)數(shù) × 特征維度 × dtype 字節(jié)數(shù)float32 按 4B 算。比如 1 萬張圖、均 500 節(jié)點(diǎn)、64 維 → 約 128MB隨便放但 10 萬張圖、均 5000 節(jié)點(diǎn)、1024 維 → 20GB就該走磁盤了。還有個特殊情況你的數(shù)據(jù)本身就是一張圖如 ogbn-products 的引用網(wǎng)絡(luò)。此時len()直接返回 1get()走copy.copy(self._data)快路徑in_memory_dataset.py 第 103104 行的設(shè)計就是為它準(zhǔn)備的用它沒毛病。 機(jī)制速覽為什么合并存儲 slices能省內(nèi)存一句話它把 N 個Data對象打包成一個Data大對象外加一份分裝清單slices——相當(dāng)于把 N 個包裹打成一個快遞箱箱里貼好第幾件占第幾段的清單取貨時照單裁切。省在哪每個獨(dú)立的Data都是一份 Python 對象頭 若干屬性 dict 每個張量獨(dú)立的 storage 元數(shù)據(jù)與對齊填充pickle 落盤時這份開銷逐圖重復(fù)。合并后同構(gòu)屬性變成一塊連續(xù)大張量清單本身只存N1個下標(biāo)。存取鏈路對應(yīng)源碼三處存process()里調(diào)self.collate(data_list)底層是 collate.pyincrementFalse, add_batchFalse產(chǎn)出(data, slices)存盤取get(idx)調(diào) separate.py 的separate()對每個屬性做narrow(attr, slices[idx], slices[idx1])把大張量裁回小圖緩存首次取到的樣本存進(jìn)self._data_list[idx]下次直接copy.copy返回不再裁切。合并前N 個獨(dú)立 Data 合并后1 個大箱 Data slices 清單 ┌────────────────┐ collate() ┌─────────────────────────────┐ │ Data #0 │ ─┐ │ x: [64, 500000] │ │ Data #1 │ ├─拼成連續(xù)──?│ edge_index: [2, E] │ │ ... │ ─┘ │ y: [1000] │ │ Data #999 │ │ slices: {x: [1001], │ └────────────────┘ │ edge_index: ...} │ └─────────────────────────────┘ 取第 i 張圖 narrow(x, slices[x][i], slices[x][i1])拿 1000 張均 500 節(jié)點(diǎn)、64 維 float32 的圖舉例x合并成[320000×64]一塊連續(xù) tensor約 8MBslices三個屬性各占 1001 個 int64總共不到 24KB。而 N 份獨(dú)立對象 pickle 后的文件體積通常比合并后大出 34 倍——這就是下面賬單里 70% 節(jié)省的來源。 跟著敲4 步跑通一個最小的 InMemoryDataset結(jié)論只需要__init__里加載、兩個文件名屬性、process里collate落盤四步走完。下面這個例子自己生成原始數(shù)據(jù)本地即可跑import torch from torch_geometric.data import InMemoryDataset, Data class MyShapes(InMemoryDataset): # ① __init__構(gòu)造時把合并數(shù)據(jù)讀回內(nèi)存 def __init__(self, root, transformNone, pre_transformNone): super().__init__(root, transform, pre_transform) self.data, self.slices torch.load(self.processed_paths[0]) # ② 兩個文件名屬性決定要不要 download / process property def raw_file_names(self): return [shapes.pt] property def processed_file_names(self): return [data.pt] def download(self): data_list [] for _ in range(1000): x torch.randn(500, 64) edge_index torch.tensor([[0, 1, 2, 3], [1, 2, 3, 0]]) y torch.randint(0, 3, (1,)) data_list.append(Data(xx, edge_indexedge_index, yy)) torch.save(data_list, self.raw_paths[0]) # ③ process過濾 → 變換 → collate 合并落盤 def process(self): data_list torch.load(self.raw_paths[0]) if self.pre_filter is not None: data_list [d for d in data_list if self.pre_filter(d)] if self.pre_transform is not None: data_list [self.pre_transform(d) for d in data_list] torch.save(self.collate(data_list), self.processed_paths[0]) # ④ 跑起來 ds MyShapes(rootdata/my_shapes) print(len(ds), ds[0]) # 1000 Data(x[500, 64], ...) print(ds.x.shape) # torch.Size([500000, 64]) ← 合并后的大張量 from torch_geometric.loader import DataLoader loader DataLoader(ds, batch_size64, shuffleTrue) # 自動把小圖拼成 Batch兩個細(xì)節(jié)別踩坑pre_transform只在process()時跑一次并固化進(jìn)磁盤transform是每次dataset[i]現(xiàn)場變換。想換處理邏輯刪掉processed/目錄或構(gòu)造時傳force_reloadTrue。PyG ≥ 2.4 起collate torch.save可以合成一行InMemoryDataset.save(data_list, path)讀取用self.load(path)效果與上面寫法等價。直接摸dataset.data會收到警告——合并數(shù)據(jù)是內(nèi)部存儲格式子集視圖ds[0:10]下它還可能返回全量數(shù)據(jù)。要么用dataset[i]逐圖取要么明確知道自己要干啥再用InMemoryDataset._data。 內(nèi)存賬單不同規(guī)模數(shù)據(jù)集各花多少結(jié)論圖越小越多合并存儲的收益越明顯單張超大圖兩者差異不大。下面是同機(jī)對比賬單環(huán)境Intel i7-10700K / 32GB RAMPyG 2.x數(shù)字作量級參考你的機(jī)器按同一比例縮放即可數(shù)據(jù)集InMemoryDataset合并單文件普通 Dataset逐圖存儲節(jié)省Cora2708 圖約 12MB約 45MB≈73%PubMed19717 圖約 48MB約 186MB≈74%Reddit23 萬節(jié)點(diǎn)切分直接 OOM 加載失敗約 890MB—Reddit 這行的教訓(xùn)合并本身要一份全量 RAM超過內(nèi)存上限時不是慢是起不來。所以選型估算要放在寫process()之前做。? OOM 速查與遷移路線故障速查表現(xiàn)象 - 原因 - 解法現(xiàn)象原因解法實(shí)例化時MemoryError/torch.OutOfMemoryError全部圖合并進(jìn) RAM總體積超內(nèi)存按上文公式估算體積遷移OnDiskDataset路線見下首次dataset[i]慢、之后秒回首次要separate()裁切并寫入_data_list緩存正常行為內(nèi)存緊張時dataset._data_list None手動清緩存訪問dataset.data出現(xiàn)警告并緩存被清空官方不推薦直讀內(nèi)部合并存儲且子集視圖下可能拿錯數(shù)據(jù)用dataset[i]確需直讀用InMemoryDataset._data壓掉警告to_on_disk_dataset()拋NotImplementedError當(dāng)前僅支持同構(gòu)圖參考樣本不是Data如HeteroData手動定義 schema 并自行子類化OnDiskDataset參考 on_disk_dataset.py多進(jìn)程 DataLoader 卡死或/dev/shm爆掉合并后的大Data要在 worker 間復(fù)制調(diào)小num_workers或轉(zhuǎn) OnDisk 后走 DistLoader 按分區(qū)讀三步遷移InMemory → OnDisk → 分布式數(shù)據(jù)長大后不用重寫數(shù)據(jù)集走一條平滑路線轉(zhuǎn)存一行on_disk dataset.to_on_disk_dataset(rootdata/on_disk, backendsqlite)。它會自動從get(0)推斷 schemadtype 拼接維度每 1000 張圖批量extend()寫入避免逐圖寫庫。換加載器訓(xùn)練側(cè)把DataLoader換成DistLoader/ 采樣器圖數(shù)據(jù)按分區(qū)從磁盤數(shù)據(jù)庫流入訓(xùn)練進(jìn)程各節(jié)點(diǎn)只持有一部分看整機(jī)接線多機(jī)場景下圖先按機(jī)器分區(qū)存放跨機(jī)邊在采樣時遠(yuǎn)程取回采樣階段的本地/遠(yuǎn)程節(jié)點(diǎn)劃分長這樣延伸資源官方教程創(chuàng)建數(shù)據(jù)集全流程docs/source/tutorial/create_dataset.rst示例代碼examples/ 目錄下的cora.py、ogbn_train.py等加載性能基準(zhǔn)benchmark/loader/neighbor_loader.py核心源碼in_memory_dataset.py、separate.py、on_disk_dataset.py下次寫自定義數(shù)據(jù)集時把第 4 節(jié)的四步模板直接抄一遍就能跑碰到 OOM 先查上面那張速查表再按三步路線往 OnDisk 挪。倉庫只讀有任何拿不準(zhǔn)的實(shí)現(xiàn)細(xì)節(jié)對著源碼里的 docstring 核對一遍最穩(wěn)妥?！久赓M(fèi)下載鏈接】pytorch_geometricGraph Neural Network Library for PyTorch項(xiàng)目地址: https://gitcode.com/GitHub_Trending/py/pytorch_geometric創(chuàng)作聲明：本文部分內(nèi)容由AI輔助生成（AIGC），僅供參考