存減重、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彈出來(lái)內(nèi)存條直接頂?shù)?98%在 PyTorch GeometricPyG里InMemoryDataset的合并存儲(chǔ) 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總體積合并后單個(gè).pt約幾百 MB幾 GB 一半 RAM超過(guò) RAM或需要多機(jī)共享圖的規(guī)模幾百幾萬(wàn)個(gè)中小圖百萬(wàn)級(jí)圖papers100m 量級(jí)或單張超大圖訪問(wèn)模式隨機(jī)取圖 DataLoader自動(dòng)合批按索引/分區(qū)點(diǎn)查訓(xùn)練形態(tài)單機(jī)、多進(jìn)程 DataLoader多機(jī)分布式DistLoader讀取速度最快純 RAM命中緩存近乎零開(kāi)銷(xiāo)每圖多一次數(shù)據(jù)庫(kù)往返略慢一個(gè)可操作的估算式圖數(shù)量 × 平均節(jié)點(diǎn)數(shù) × 特征維度 × dtype 字節(jié)數(shù)float32 按 4B 算。比如 1 萬(wàn)張圖、均 500 節(jié)點(diǎn)、64 維 → 約 128MB隨便放但 10 萬(wàn)張圖、均 5000 節(jié)點(diǎn)、1024 維 → 20GB就該走磁盤(pán)了。還有個(gè)特殊情況你的數(shù)據(jù)本身就是一張圖如 ogbn-products 的引用網(wǎng)絡(luò)。此時(shí)len()直接返回 1get()走copy.copy(self._data)快路徑in_memory_dataset.py 第 103104 行的設(shè)計(jì)就是為它準(zhǔn)備的用它沒(méi)毛病。 機(jī)制速覽為什么合并存儲(chǔ) slices能省內(nèi)存一句話它把 N 個(gè)Data對(duì)象打包成一個(gè)Data大對(duì)象外加一份分裝清單slices——相當(dāng)于把 N 個(gè)包裹打成一個(gè)快遞箱箱里貼好第幾件占第幾段的清單取貨時(shí)照單裁切。省在哪每個(gè)獨(dú)立的Data都是一份 Python 對(duì)象頭 若干屬性 dict 每個(gè)張量獨(dú)立的 storage 元數(shù)據(jù)與對(duì)齊填充pickle 落盤(pán)時(shí)這份開(kāi)銷(xiāo)逐圖重復(fù)。合并后同構(gòu)屬性變成一塊連續(xù)大張量清單本身只存N1個(gè)下標(biāo)。存取鏈路對(duì)應(yīng)源碼三處存process()里調(diào)self.collate(data_list)底層是 collate.pyincrementFalse, add_batchFalse產(chǎn)出(data, slices)存盤(pán)取get(idx)調(diào) separate.py 的separate()對(duì)每個(gè)屬性做narrow(attr, slices[idx], slices[idx1])把大張量裁回小圖緩存首次取到的樣本存進(jìn)self._data_list[idx]下次直接copy.copy返回不再裁切。合并前N 個(gè)獨(dú)立 Data 合并后1 個(gè)大箱 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三個(gè)屬性各占 1001 個(gè) int64總共不到 24KB。而 N 份獨(dú)立對(duì)象 pickle 后的文件體積通常比合并后大出 34 倍——這就是下面賬單里 70% 節(jié)省的來(lái)源。 跟著敲4 步跑通一個(gè)最小的 InMemoryDataset結(jié)論只需要__init__里加載、兩個(gè)文件名屬性、process里collate落盤(pán)四步走完。下面這個(gè)例子自己生成原始數(shù)據(jù)本地即可跑import torch from torch_geometric.data import InMemoryDataset, Data class MyShapes(InMemoryDataset): # ① __init__構(gòu)造時(shí)把合并數(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]) # ② 兩個(gè)文件名屬性決定要不要 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過(guò)濾 → 變換 → collate 合并落盤(pán) 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]) # ④ 跑起來(lái) 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) # 自動(dòng)把小圖拼成 Batch兩個(gè)細(xì)節(jié)別踩坑pre_transform只在process()時(shí)跑一次并固化進(jìn)磁盤(pán)transform是每次dataset[i]現(xiàn)場(chǎng)變換。想換處理邏輯刪掉processed/目錄或構(gòu)造時(shí)傳force_reloadTrue。PyG ≥ 2.4 起collate torch.save可以合成一行InMemoryDataset.save(data_list, path)讀取用self.load(path)效果與上面寫(xiě)法等價(jià)。直接摸dataset.data會(huì)收到警告——合并數(shù)據(jù)是內(nèi)部存儲(chǔ)格式子集視圖ds[0:10]下它還可能返回全量數(shù)據(jù)。要么用dataset[i]逐圖取要么明確知道自己要干啥再用InMemoryDataset._data。 內(nèi)存賬單不同規(guī)模數(shù)據(jù)集各花多少結(jié)論圖越小越多合并存儲(chǔ)的收益越明顯單張超大圖兩者差異不大。下面是同機(jī)對(duì)比賬單環(huán)境Intel i7-10700K / 32GB RAMPyG 2.x數(shù)字作量級(jí)參考你的機(jī)器按同一比例縮放即可數(shù)據(jù)集InMemoryDataset合并單文件普通 Dataset逐圖存儲(chǔ)節(jié)省Cora2708 圖約 12MB約 45MB≈73%PubMed19717 圖約 48MB約 186MB≈74%Reddit23 萬(wàn)節(jié)點(diǎn)切分直接 OOM 加載失敗約 890MB—Reddit 這行的教訓(xùn)合并本身要一份全量 RAM超過(guò)內(nèi)存上限時(shí)不是慢是起不來(lái)。所以選型估算要放在寫(xiě)process()之前做。? OOM 速查與遷移路線故障速查表現(xiàn)象 - 原因 - 解法現(xiàn)象原因解法實(shí)例化時(shí)MemoryError/torch.OutOfMemoryError全部圖合并進(jìn) RAM總體積超內(nèi)存按上文公式估算體積遷移OnDiskDataset路線見(jiàn)下首次dataset[i]慢、之后秒回首次要separate()裁切并寫(xiě)入_data_list緩存正常行為內(nèi)存緊張時(shí)dataset._data_list None手動(dòng)清緩存訪問(wèn)dataset.data出現(xiàn)警告并緩存被清空官方不推薦直讀內(nèi)部合并存儲(chǔ)且子集視圖下可能拿錯(cuò)數(shù)據(jù)用dataset[i]確需直讀用InMemoryDataset._data壓掉警告to_on_disk_dataset()拋NotImplementedError當(dāng)前僅支持同構(gòu)圖參考樣本不是Data如HeteroData手動(dòng)定義 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ù)長(zhǎng)大后不用重寫(xiě)數(shù)據(jù)集走一條平滑路線轉(zhuǎn)存一行on_disk dataset.to_on_disk_dataset(rootdata/on_disk, backendsqlite)。它會(huì)自動(dòng)從get(0)推斷 schemadtype 拼接維度每 1000 張圖批量extend()寫(xiě)入避免逐圖寫(xiě)庫(kù)。換加載器訓(xùn)練側(cè)把DataLoader換成DistLoader/ 采樣器圖數(shù)據(jù)按分區(qū)從磁盤(pán)數(shù)據(jù)庫(kù)流入訓(xùn)練進(jìn)程各節(jié)點(diǎn)只持有一部分看整機(jī)接線多機(jī)場(chǎng)景下圖先按機(jī)器分區(qū)存放跨機(jī)邊在采樣時(shí)遠(yuǎn)程取回采樣階段的本地/遠(yuǎn)程節(jié)點(diǎn)劃分長(zhǎng)這樣延伸資源官方教程創(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下次寫(xiě)自定義數(shù)據(jù)集時(shí)把第 4 節(jié)的四步模板直接抄一遍就能跑碰到 OOM 先查上面那張速查表再按三步路線往 OnDisk 挪。倉(cāng)庫(kù)只讀有任何拿不準(zhǔn)的實(shí)現(xiàn)細(xì)節(jié)對(duì)著源碼里的 docstring 核對(duì)一遍最穩(wěn)妥。【免費(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彈出來(lái)內(nèi)存條直接頂?shù)?98%在 PyTorch GeometricPyG里InMemoryDataset的合并存儲(chǔ) 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總體積合并后單個(gè).pt約幾百 MB幾 GB 一半 RAM超過(guò) RAM或需要多機(jī)共享圖的規(guī)模幾百幾萬(wàn)個(gè)中小圖百萬(wàn)級(jí)圖papers100m 量級(jí)或單張超大圖訪問(wèn)模式隨機(jī)取圖 DataLoader自動(dòng)合批按索引/分區(qū)點(diǎn)查訓(xùn)練形態(tài)單機(jī)、多進(jìn)程 DataLoader多機(jī)分布式DistLoader讀取速度最快純 RAM命中緩存近乎零開(kāi)銷(xiāo)每圖多一次數(shù)據(jù)庫(kù)往返略慢一個(gè)可操作的估算式圖數(shù)量 × 平均節(jié)點(diǎn)數(shù) × 特征維度 × dtype 字節(jié)數(shù)float32 按 4B 算。比如 1 萬(wàn)張圖、均 500 節(jié)點(diǎn)、64 維 → 約 128MB隨便放但 10 萬(wàn)張圖、均 5000 節(jié)點(diǎn)、1024 維 → 20GB就該走磁盤(pán)了。還有個(gè)特殊情況你的數(shù)據(jù)本身就是一張圖如 ogbn-products 的引用網(wǎng)絡(luò)。此時(shí)len()直接返回 1get()走copy.copy(self._data)快路徑in_memory_dataset.py 第 103104 行的設(shè)計(jì)就是為它準(zhǔn)備的用它沒(méi)毛病。 機(jī)制速覽為什么合并存儲(chǔ) slices能省內(nèi)存一句話它把 N 個(gè)Data對(duì)象打包成一個(gè)Data大對(duì)象外加一份分裝清單slices——相當(dāng)于把 N 個(gè)包裹打成一個(gè)快遞箱箱里貼好第幾件占第幾段的清單取貨時(shí)照單裁切。省在哪每個(gè)獨(dú)立的Data都是一份 Python 對(duì)象頭 若干屬性 dict 每個(gè)張量獨(dú)立的 storage 元數(shù)據(jù)與對(duì)齊填充pickle 落盤(pán)時(shí)這份開(kāi)銷(xiāo)逐圖重復(fù)。合并后同構(gòu)屬性變成一塊連續(xù)大張量清單本身只存N1個(gè)下標(biāo)。存取鏈路對(duì)應(yīng)源碼三處存process()里調(diào)self.collate(data_list)底層是 collate.pyincrementFalse, add_batchFalse產(chǎn)出(data, slices)存盤(pán)取get(idx)調(diào) separate.py 的separate()對(duì)每個(gè)屬性做narrow(attr, slices[idx], slices[idx1])把大張量裁回小圖緩存首次取到的樣本存進(jìn)self._data_list[idx]下次直接copy.copy返回不再裁切。合并前N 個(gè)獨(dú)立 Data 合并后1 個(gè)大箱 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三個(gè)屬性各占 1001 個(gè) int64總共不到 24KB。而 N 份獨(dú)立對(duì)象 pickle 后的文件體積通常比合并后大出 34 倍——這就是下面賬單里 70% 節(jié)省的來(lái)源。 跟著敲4 步跑通一個(gè)最小的 InMemoryDataset結(jié)論只需要__init__里加載、兩個(gè)文件名屬性、process里collate落盤(pán)四步走完。下面這個(gè)例子自己生成原始數(shù)據(jù)本地即可跑import torch from torch_geometric.data import InMemoryDataset, Data class MyShapes(InMemoryDataset): # ① __init__構(gòu)造時(shí)把合并數(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]) # ② 兩個(gè)文件名屬性決定要不要 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過(guò)濾 → 變換 → collate 合并落盤(pán) 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]) # ④ 跑起來(lái) 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) # 自動(dòng)把小圖拼成 Batch兩個(gè)細(xì)節(jié)別踩坑pre_transform只在process()時(shí)跑一次并固化進(jìn)磁盤(pán)transform是每次dataset[i]現(xiàn)場(chǎng)變換。想換處理邏輯刪掉processed/目錄或構(gòu)造時(shí)傳force_reloadTrue。PyG ≥ 2.4 起collate torch.save可以合成一行InMemoryDataset.save(data_list, path)讀取用self.load(path)效果與上面寫(xiě)法等價(jià)。直接摸dataset.data會(huì)收到警告——合并數(shù)據(jù)是內(nèi)部存儲(chǔ)格式子集視圖ds[0:10]下它還可能返回全量數(shù)據(jù)。要么用dataset[i]逐圖取要么明確知道自己要干啥再用InMemoryDataset._data。 內(nèi)存賬單不同規(guī)模數(shù)據(jù)集各花多少結(jié)論圖越小越多合并存儲(chǔ)的收益越明顯單張超大圖兩者差異不大。下面是同機(jī)對(duì)比賬單環(huán)境Intel i7-10700K / 32GB RAMPyG 2.x數(shù)字作量級(jí)參考你的機(jī)器按同一比例縮放即可數(shù)據(jù)集InMemoryDataset合并單文件普通 Dataset逐圖存儲(chǔ)節(jié)省Cora2708 圖約 12MB約 45MB≈73%PubMed19717 圖約 48MB約 186MB≈74%Reddit23 萬(wàn)節(jié)點(diǎn)切分直接 OOM 加載失敗約 890MB—Reddit 這行的教訓(xùn)合并本身要一份全量 RAM超過(guò)內(nèi)存上限時(shí)不是慢是起不來(lái)。所以選型估算要放在寫(xiě)process()之前做。? OOM 速查與遷移路線故障速查表現(xiàn)象 - 原因 - 解法現(xiàn)象原因解法實(shí)例化時(shí)MemoryError/torch.OutOfMemoryError全部圖合并進(jìn) RAM總體積超內(nèi)存按上文公式估算體積遷移OnDiskDataset路線見(jiàn)下首次dataset[i]慢、之后秒回首次要separate()裁切并寫(xiě)入_data_list緩存正常行為內(nèi)存緊張時(shí)dataset._data_list None手動(dòng)清緩存訪問(wèn)dataset.data出現(xiàn)警告并緩存被清空官方不推薦直讀內(nèi)部合并存儲(chǔ)且子集視圖下可能拿錯(cuò)數(shù)據(jù)用dataset[i]確需直讀用InMemoryDataset._data壓掉警告to_on_disk_dataset()拋NotImplementedError當(dāng)前僅支持同構(gòu)圖參考樣本不是Data如HeteroData手動(dòng)定義 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ù)長(zhǎng)大后不用重寫(xiě)數(shù)據(jù)集走一條平滑路線轉(zhuǎn)存一行on_disk dataset.to_on_disk_dataset(rootdata/on_disk, backendsqlite)。它會(huì)自動(dòng)從get(0)推斷 schemadtype 拼接維度每 1000 張圖批量extend()寫(xiě)入避免逐圖寫(xiě)庫(kù)。換加載器訓(xùn)練側(cè)把DataLoader換成DistLoader/ 采樣器圖數(shù)據(jù)按分區(qū)從磁盤(pán)數(shù)據(jù)庫(kù)流入訓(xùn)練進(jìn)程各節(jié)點(diǎn)只持有一部分看整機(jī)接線多機(jī)場(chǎng)景下圖先按機(jī)器分區(qū)存放跨機(jī)邊在采樣時(shí)遠(yuǎn)程取回采樣階段的本地/遠(yuǎn)程節(jié)點(diǎn)劃分長(zhǎng)這樣延伸資源官方教程創(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下次寫(xiě)自定義數(shù)據(jù)集時(shí)把第 4 節(jié)的四步模板直接抄一遍就能跑碰到 OOM 先查上面那張速查表再按三步路線往 OnDisk 挪。倉(cāng)庫(kù)只讀有任何拿不準(zhǔn)的實(shí)現(xiàn)細(xì)節(jié)對(duì)著源碼里的 docstring 核對(duì)一遍最穩(wěn)妥。【免費(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），僅供參考