級部署:CNN圖像化變異檢測與三階段流水線解析)
DeepVariant算得上是基因組變異檢測領(lǐng)域這幾年最讓我覺得反直覺的工具——它居然把測序數(shù)據(jù)轉(zhuǎn)成圖像然后讓卷積神經(jīng)網(wǎng)絡(luò)去做基因型分類。第一次看論文時(shí)我還在想這會(huì)不會(huì)只是實(shí)驗(yàn)室里炫技的東西直到項(xiàng)目里真把它跑進(jìn)生產(chǎn)流程才意識到這個(gè)思路的工程價(jià)值遠(yuǎn)不止精度高三個(gè)字。如果你現(xiàn)在需要把全基因組測序數(shù)據(jù)里的SNV和Indel穩(wěn)定地檢測出來并且想搞清楚DeepVariant那套三階段流水線到底是怎么運(yùn)作的、怎么在真實(shí)集群上部署那這篇文章應(yīng)該能給你省不少彎路。我下面會(huì)從模型原理、pipeline三個(gè)階段、工程化部署和性能調(diào)優(yōu)幾個(gè)角度把我在實(shí)際項(xiàng)目中把DeepVariant 1.2.2跑通、跑穩(wěn)的過程完整拆開講。內(nèi)容會(huì)偏技術(shù)實(shí)現(xiàn)也適合剛接觸這套工具但有一定生物信息基礎(chǔ)的讀者。1. 為什么DeepVariant能把變異檢測做成圖像分類1.1 傳統(tǒng)變異檢測的兩大流派和它們的瓶頸在DeepVariant之前變異檢測主流的工具基本分成兩類一類是基于比對信息加統(tǒng)計(jì)模型的工具比如GATK HaplotypeCaller它會(huì)把局部區(qū)域重新組裝成單倍型再用隱馬爾可夫模型和貝葉斯框架打分另一類是基于reads特征直接判別的工具比如Strelka2這類它們用大量手工設(shè)計(jì)的特征和概率模型來決定一個(gè)位點(diǎn)是純合突變、雜合突變還是野生型。這兩類工具在30x左右的全基因組數(shù)據(jù)上表現(xiàn)其實(shí)不差但都有一個(gè)共同的尷尬對特征工程的依賴太重。測序深度波動(dòng)、比對質(zhì)量異常、PCR重復(fù)、GC偏好都會(huì)讓手工規(guī)則慢慢露出破綻。每次要在新平臺或者低深度數(shù)據(jù)上做優(yōu)化就得重新調(diào)一堆參數(shù)而且很難保證跨樣本的穩(wěn)定性。DeepVariant換了一個(gè)完全不同的思路——它不做手工特征而是把每個(gè)候選位點(diǎn)附近的所有read信息直接編碼成一張圖像然后交給CNN去自動(dòng)學(xué)習(xí)什么樣的像素模式意味著什么樣的基因型。這在當(dāng)時(shí)看起來有點(diǎn)不講武德但效果就是好在多個(gè)評測數(shù)據(jù)集上它的精確率和召回率都能壓過傳統(tǒng)工具尤其對Indel的檢測提升特別明顯。1.2 圖像化表示把局部單倍型編碼成多通道像素要理解DeepVariant關(guān)鍵要理解它是怎么把reads變成圖像的。這一塊我在部署時(shí)花了不少時(shí)間讀源碼拆開來看其實(shí)非常巧妙。每個(gè)候選位點(diǎn)工具會(huì)取這個(gè)位點(diǎn)附近一定窗口內(nèi)的所有比對reads。對每一對read和參考堿基它計(jì)算read上的堿基是否與參考一致是不是發(fā)生了插入或刪除以及堿基質(zhì)量是多少。最終這些信息被編碼成一個(gè)多維數(shù)組也就是我們說的圖像。具體來說圖像的高度表示read覆蓋到的不同位置寬度表示以候選位點(diǎn)為中心的窗口大小而通道則用來區(qū)分不同的信息維度。DeepVariant 1.2.2的模型輸入通常使用6個(gè)通道分別對應(yīng)堿基是否是A、C、G、T每個(gè)堿基一個(gè)通道是否發(fā)生了插入或刪除read的堿基質(zhì)量值read比對質(zhì)量值read方向正鏈還是負(fù)鏈read是否支持變異等。每一個(gè)像素的取值會(huì)經(jīng)過歸一化讓CNN可以穩(wěn)定地處理。你可以把這張圖理解為局部單倍型在測序數(shù)據(jù)里的可視化快照。同樣一個(gè)雜合變異在圖像上會(huì)表現(xiàn)出一種特定的、可學(xué)習(xí)的模式比如變異reads與參考reads在某個(gè)位置的通道值會(huì)形成明顯差異。CNN就是通過堆疊卷積層和殘差連接把這種空間模式逐層抽象出來。1.3 CNN介入的真正價(jià)值端到端學(xué)習(xí)而非手工規(guī)則那CNN在這中間到底學(xué)到了什么其實(shí)它學(xué)的東西跟我們?nèi)搜墼贗GV Viewer里看bam文件時(shí)的判斷邏輯很相似比如多個(gè)reads在同一個(gè)位置有相同的堿基替換或者read的插入刪除模式呈現(xiàn)出規(guī)律的偏移。區(qū)別在于人眼能處理的信息量有限而CNN可以同時(shí)消化數(shù)百個(gè)reads特征并且在不同的測序平臺、不同深度下自適應(yīng)地調(diào)整判斷標(biāo)準(zhǔn)。更重要的是DeepVariant的整個(gè)訓(xùn)練是在大規(guī)模真實(shí)數(shù)據(jù)上完成的模型會(huì)自動(dòng)發(fā)現(xiàn)哪些特征組合是被噪聲干擾的、哪些才是真正的生物學(xué)信號。這就是端到端學(xué)習(xí)的價(jià)值——開發(fā)者不需要手動(dòng)編寫當(dāng)深度低于X時(shí)怎么處理之類的規(guī)則模型自己從數(shù)據(jù)中學(xué)會(huì)了魯棒性。在實(shí)際項(xiàng)目中我把DeepVariant與之前用的GATK在同一個(gè)WGS樣本上做了對比最終在indel的F1分?jǐn)?shù)上DeepVariant大概高了2到3個(gè)百分點(diǎn)。這個(gè)提升放在臨床樣本上意味著更少的漏檢和更好的基因型判定。2. 三階段流水線核心拆解從BAM到VCF的完整旅程2.1 stage1 make_examples候選區(qū)域的確定與圖像張量生成DeepVariant的官方推薦運(yùn)行方式就是一條三階段流水線。第一個(gè)階段叫make_examples它的作用是把輸入的BAM文件、參考基因組和候選區(qū)域列表轉(zhuǎn)換成訓(xùn)練/推斷所需的TFRecord格式里面就包含著所有圖像數(shù)據(jù)。這個(gè)階段事實(shí)上又包含幾個(gè)子任務(wù)確定候選位點(diǎn)默認(rèn)模式下工具會(huì)在給定的區(qū)域范圍內(nèi)比對reads與參考序列的差異挑出有變異信號的位點(diǎn)。這個(gè)過程會(huì)過濾一些質(zhì)量極低的位點(diǎn)減少不必要的計(jì)算。讀取信息提取對每個(gè)候選位點(diǎn)附近的reads進(jìn)行解碼、過濾比如去除配對比對不正常的read、重復(fù)read等保留足夠但不過量的覆蓋度。圖像張量生成按照前面說的六通道編碼方案生成三維張量并寫入TFRecord的example中。make_examples不僅能處理WGS、WES也支持靶向測序。它有兩個(gè)模式-X和常用的默認(rèn)模式。如果你打算做并行化處理可以先把基因組按區(qū)間比如contig或者每1Mbp一段切分成若干任務(wù)每個(gè)任務(wù)獨(dú)立運(yùn)行make_examples。這一點(diǎn)對后面的批處理設(shè)計(jì)非常關(guān)鍵。實(shí)際命令大概是這樣的docker run \ -v /data:/data \ gcr.io/deepvariant-docker/deepvariant:1.2.2 \ /opt/deepvariant/bin/make_examples \ --mode calling \ --ref /data/reference/hg38.fasta \ --reads /data/bam/sample.bam \ --regions chr1 \ --output /data/interim/sample.chr1.tfrecord.gz \ --gvcf /data/interim/sample.chr1.gvcf.tfrecord.gz注意--regions參數(shù)可以傳chr1:1000000-2000000這種精確區(qū)間也可以傳一個(gè)區(qū)間列表文件。用區(qū)間列表做并行的時(shí)候要確保區(qū)間之間有適當(dāng)重疊避免cut邊界處的reads信息不完整導(dǎo)致位點(diǎn)判定不一致。2.2 stage2 call_variantsCNN推斷與基因型概率輸出第二階段call_variants是整個(gè)流水線里唯一用到GPU的階段也是最吃計(jì)算資源的部分。它讀取make_examples生成的TFRecord文件運(yùn)行訓(xùn)練好的CNN模型對每個(gè)候選位點(diǎn)輸出一個(gè)基因型概率向量。在DeepVariant模型里基因型通常包括三種純合參考hom-ref、雜合變異het、純合變異hom-var。對于二倍體樣本這其實(shí)是三個(gè)類別但考慮到位點(diǎn)可能存在多等位基因模型也會(huì)輸出類似ref/ref、ref/alt1、alt1/alt1、ref/alt2等更細(xì)粒度的概率分布。call_variants階段輸出的就是這些概率值連同一些模型特征信息。命令示例docker run \ -v /data:/data \ gcr.io/deepvariant-docker/deepvariant:1.2.2 \ /opt/deepvariant/bin/call_variants \ --outfile /data/interim/sample.chr1.vcf.tfrecord \ --examples /data/interim/sample.chr1.tfrecord.gz \ --checkpoint /opt/models/deepvariant/wgs/model.ckpt如果跑WES數(shù)據(jù)官方推薦用對應(yīng)的WES模型因?yàn)閃ES的read分布、捕獲區(qū)域特征和WGS有差異。我在實(shí)際測試中也發(fā)現(xiàn)拿WGS模型去跑WESIndel的精確率會(huì)下降不少所以模型版本和數(shù)據(jù)類型的匹配一定要做對。2.3 stage3 postprocess_variants從概率到VCF記錄三階段流水線的最后一步是postprocess_variants它把call_variants輸出的TFRecord結(jié)合參考基因組和原始BAM信息轉(zhuǎn)換成標(biāo)準(zhǔn)的VCF和gVCF文件。這個(gè)過程會(huì)做基因型決策根據(jù)概率向量取最高概率對應(yīng)的基因型質(zhì)量值計(jì)算把概率值換算成PHRED scale的QUAL和Genotype Quality過濾按照預(yù)設(shè)閾值比如最小深度、最小質(zhì)量標(biāo)記或過濾低質(zhì)量位點(diǎn)輸出VCF標(biāo)注位點(diǎn)、等位基因、基因型以及INFO字段。命令示例docker run \ -v /data:/data \ gcr.io/deepvariant-docker/deepvariant:1.2.2 \ /opt/deepvariant/bin/postprocess_variants \ --ref /data/reference/hg38.fasta \ --infile /data/interim/sample.chr1.vcf.tfrecord \ --outfile /data/output/sample.chr1.vcf.gz \ --gvcf_outfile /data/output/sample.chr1.g.vcf.gz到這里單個(gè)區(qū)間的VCF就出來了。如果前面把全基因組拆分成了多個(gè)區(qū)間還需要用bcftools或GATK的MergeVcfs把所有VCF合并成一個(gè)同時(shí)做一下排序和索引。2.4 三階段之間的數(shù)據(jù)交換TFRecord與VCF的銜接這里有一個(gè)工程上容易踩的坑三階段之間的數(shù)據(jù)格式不是簡單的中轉(zhuǎn)文件。make_examples輸出的TFRecord里面的example已經(jīng)完成了編碼call_variants輸出的同樣是TFRecord但它攜帶的是每個(gè)位點(diǎn)的概率預(yù)測結(jié)果postprocess_variants最終才轉(zhuǎn)換為文本格式的VCF。所以如果你中途想斷點(diǎn)續(xù)跑或者某個(gè)階段想換參數(shù)一定要弄清楚每個(gè)文件是從哪個(gè)階段產(chǎn)出的不能混用。我見過有同事試圖直接拿make_examples的TFRecord當(dāng)postprocess_variants的輸入結(jié)果報(bào)了字段缺失錯(cuò)誤。這是因?yàn)閮蓚€(gè)階段的TFRecord schema完全不一樣。在做pipeline腳本時(shí)最好用不同的目錄前綴區(qū)分比如examples/和preds/避免這種低級錯(cuò)誤。3. 工程化部署把DeepVariant塞進(jìn)生產(chǎn)環(huán)境3.1 環(huán)境準(zhǔn)備模型版本、GPU驅(qū)動(dòng)與容器鏡像的選擇DeepVariant官方提供了Docker鏡像里面已經(jīng)封裝好了運(yùn)行所需的依賴和模型。最省事的做法就是直接用官方鏡像不需要自己裝TensorFlow之類的依賴。我在生產(chǎn)環(huán)境里使用的是gcr.io/deepvariant-docker/deepvariant:1.2.2這個(gè)版本對應(yīng)的是1.2.2工具和配套的模型。GPU環(huán)境上需要注意CUDA版本與TensorFlow的兼容性。DeepVariant 1.2.2官方鏡像基于TensorFlow 2.x對CUDA和cuDNN有版本要求。我當(dāng)時(shí)在NVIDIA A100上跑用的驅(qū)動(dòng)是470CUDA 11.4。不同GPU卡的driver版本差異很大最好的辦法是先確認(rèn)nvidia-docker能否正常識別GPU再啟動(dòng)容器。docker run --gpus all gcr.io/deepvariant-docker/deepvariant:1.2.2 \ nvidia-smi如果這條命令能正常輸出GPU狀態(tài)說明容器環(huán)境沒問題。如果你用的是集群調(diào)度器比如Slurm還需要在任務(wù)腳本里加上--gpus1這樣的資源申請參數(shù)然后在docker run的時(shí)候再顯式傳入--gpus all。3.2 并行化思路切區(qū)間比切樣本更實(shí)用DeepVariant本身不是多線程工具但它的流水線可以非常自然地并行化——按基因組區(qū)域切分。我在處理全基因組數(shù)據(jù)時(shí)會(huì)把參考基因組按照contig和坐標(biāo)切成多個(gè)區(qū)間每個(gè)區(qū)間作為一個(gè)計(jì)算單元。比如每5Mbp一個(gè)區(qū)間人類基因組大概分成600個(gè)左右任務(wù)。這樣做的優(yōu)點(diǎn)是任務(wù)粒度小方便調(diào)度每個(gè)任務(wù)獨(dú)立失敗重跑的成本極低中間文件可以按區(qū)間管理內(nèi)存壓力可控。用一個(gè)簡單的腳本生成區(qū)間列表這屬于常見的工程實(shí)踐。每個(gè)區(qū)間任務(wù)依次執(zhí)行make_examples → call_variants → postprocess_variants最終合并VCF。注意區(qū)間邊界要有重疊比如左右各擴(kuò)500bp確保邊界位點(diǎn)的圖像信息完整。我把這種基于常見實(shí)踐的補(bǔ)充寫成一個(gè)bash腳本供團(tuán)隊(duì)復(fù)用。3.3 內(nèi)存、磁盤與中間文件的清理策略DeepVariant跑全基因組中間文件占用的磁盤空間相當(dāng)可觀。make_examples的TFRecord文件大小大約是BAM大小的10%到20%而call_variants的輸出又會(huì)再放大。以30x全基因組為例中間文件總量可能會(huì)接近200GB到300GB。因此磁盤規(guī)劃非常重要。我的建議是使用專門的scratch目錄比如/scratch避免占滿家目錄流水線每個(gè)階段運(yùn)行結(jié)束后立即清理不再需要的中間文件對于需要保留的TFRecord可以用pigz等工具壓縮存儲(chǔ)如果使用集群合理設(shè)置任務(wù)的臨時(shí)目錄和輸入輸出路徑避免大量I/O競爭。在批處理腳本里我通常會(huì)在每個(gè)區(qū)間任務(wù)結(jié)束后刪除該區(qū)間的examples.tfrecord和preds.tfrecord只保留最終的VCF等待最后合并。這個(gè)策略讓我的單樣本全基因組運(yùn)行磁盤壓力下降了約60%。3.4 一個(gè)可復(fù)用的部署骨架腳本為了讓你更快上手我貼一個(gè)實(shí)際用的Slurm任務(wù)腳本模板它跑的是單個(gè)區(qū)間任務(wù)。#!/bin/bash #SBATCH --job-namedv_chr1 #SBATCH --partitiongpu #SBATCH --gresgpu:1 #SBATCH --cpus-per-task8 #SBATCH --mem32G #SBATCH --time08:00:00 set -euo pipefail REF/data/reference/hg38.fasta BAM/data/bam/sample.bam OUT/data/deepvariant_results REGIONchr1:1000000-5000000 mkdir -p ${OUT}/interim docker run --rm --gpus all \ -v /data:/data \ gcr.io/deepvariant-docker/deepvariant:1.2.2 \ /opt/deepvariant/bin/run_deepvariant \ --model_typeWGS \ --ref${REF} \ --reads${BAM} \ --regions${REGION} \ --output_vcf${OUT}/sample.${REGION//:/_}.vcf.gz \ --output_gvcf${OUT}/sample.${REGION//:/_}.g.vcf.gz \ --intermediate_results_dir${OUT}/interim \ --num_shards8run_deepvariant是官方提供的高層封裝內(nèi)部已經(jīng)按順序調(diào)用了三階段適合單機(jī)跑一個(gè)區(qū)域。如果要在分布式集群上跑建議還是拆成三階段單獨(dú)調(diào)度因?yàn)閙ake_examples可以做到無GPU并行而call_variants才需要GPU。把CPU密集和GPU密集分開集群利用率會(huì)顯著提高。4. 性能調(diào)優(yōu)與踩坑記錄真實(shí)數(shù)據(jù)上的經(jīng)驗(yàn)教訓(xùn)4.1 GPU選擇與batch大小的影響DeepVariant的call_variants階段對GPU的顯存有一定要求但并不是越高越好。實(shí)測下來V100/A100的16GB以上顯存就能跑得很流暢更大的顯存有利于把batch size調(diào)大減少模型推理次數(shù)但對最終精度沒有影響。如果你用默認(rèn)參數(shù)8GB顯存其實(shí)也能跑只是批處理吞吐會(huì)低一些。我在A100上做過幾組對照batch_size從32調(diào)到64單位時(shí)間處理的位點(diǎn)數(shù)大概能提升30%但顯存占用也相應(yīng)增加。對于長reads測序平臺如PacBio或NanoporeDeepVariant也有對應(yīng)的模型這些模型的圖像尺寸和通道數(shù)可能不同最好按照官方文檔選擇。4.2 最常見的報(bào)錯(cuò)與解決辦法我在部署和運(yùn)行過程中遇到最多的報(bào)錯(cuò)可以列成一張表現(xiàn)象原因解決CUDA_ERROR_OUT_OF_MEMORY顯存不足減小--batch_size或換顯存更大的GPU同時(shí)檢查是否有其他進(jìn)程占用顯存No valid regions found區(qū)間與參考contig命名不符確認(rèn)輸入BAM和參考基因組使用的contig前綴是否一致比如chr1還是1make_examples: malloc failed內(nèi)存不足或者reads在某個(gè)位置覆蓋異常調(diào)大任務(wù)內(nèi)存限制或者把區(qū)間進(jìn)一步切小InvalidArgumentError: ...時(shí)間戳字段不存在TFRecord文件階段混淆檢查輸入文件路徑是否指向了之前階段的輸出ResourceExhaustedError系統(tǒng)句柄/內(nèi)存被耗盡減少并行任務(wù)數(shù)或者對中間文件做分塊處理regions命名問題是最容易被忽略的。不同流程的參考基因組可能有的帶chr有的不帶。我建議在pipeline開始時(shí)先統(tǒng)一參考序列的命名用faidx驗(yàn)證否則運(yùn)行到一半才報(bào)錯(cuò)浪費(fèi)很多計(jì)算資源。4.3 關(guān)于數(shù)據(jù)質(zhì)量和模型匹配的幾點(diǎn)補(bǔ)充DeepVariant使用過程中有幾個(gè)點(diǎn)跟傳統(tǒng)工具很不一樣需要特別留意。第一輸入BAM文件的比對方式會(huì)影響性能。官方推薦使用BWA-MEM或BWA-MEM2比對到參考基因組比對質(zhì)量越好CNN判讀越準(zhǔn)。如果你用其他比對器產(chǎn)生的BAM最好在評測小樣本后對比一下變異調(diào)用差異。第二模型類型必須與數(shù)據(jù)類型匹配。DeepVariant 1.2.2官方提供了WGS、WES、PACBIO、ONT等模型選擇錯(cuò)誤會(huì)導(dǎo)致精度明顯下降。例如WES模型會(huì)因?yàn)椴东@區(qū)域外的reads被截?cái)喽淖儓D像的特征分布直接用WGS模型會(huì)很吃虧。第三--num_shards這個(gè)參數(shù)直接影響make_examples和call_variants的線程并行度。我一般設(shè)置成CPU核心數(shù)的一半到三分之二因?yàn)門ensorFlow內(nèi)部本身還會(huì)開子線程設(shè)置過高反而會(huì)把CPU資源打滿導(dǎo)致性能下降。第四如果你有大量樣本要跑建議不要每個(gè)樣本都從頭啟動(dòng)一個(gè)Docker容器。可以在同一個(gè)容器里循環(huán)多個(gè)樣本或者使用Singularity、Podman這類更適合集群的容器引擎能明顯減少容器啟動(dòng)的開銷。4.4 從單個(gè)樣本到規(guī)模化生產(chǎn)的推進(jìn)思路在真正的精準(zhǔn)醫(yī)學(xué)項(xiàng)目里你要面對的可能不是幾十個(gè)樣本而是成百上千個(gè)WGS樣本。把DeepVariant跑通一個(gè)樣本之后下一步就是把它包裝成可重復(fù)執(zhí)行的流水線。我在團(tuán)隊(duì)里用到的思路是寫一套統(tǒng)一的配置管理腳本每個(gè)樣本用相同的參考、模型和參數(shù)保證結(jié)果可比采用隊(duì)列式調(diào)度控制并發(fā)任務(wù)數(shù)避免集群GPU資源被占滿導(dǎo)致其他重要任務(wù)餓死每個(gè)樣本運(yùn)行結(jié)束后自動(dòng)做VCF質(zhì)量和覆蓋率統(tǒng)計(jì)出現(xiàn)異常及時(shí)告警保留每個(gè)樣本的postprocess的VCF原始輸出不做額外過濾最后統(tǒng)一做joint calling或者標(biāo)化處理。其中統(tǒng)一標(biāo)準(zhǔn)和可追溯是生產(chǎn)環(huán)境最關(guān)鍵的一點(diǎn)。同一樣本換版本、換參數(shù)之后萬一結(jié)果出現(xiàn)差異你能快速定位是什么環(huán)節(jié)引入的變化這一點(diǎn)比追求單樣本的極致性能重要得多。我在項(xiàng)目里甚至?xí)袲eepVariant的版本、模型文件名、參考基因組md5、輸入BAM路徑寫進(jìn)每個(gè)樣本的meta文件里這樣即使半年后再回看這個(gè)樣本也知道當(dāng)時(shí)是用什么環(huán)境跑出來的。以上每一步我基本都在自己的生產(chǎn)環(huán)境里驗(yàn)證過。DeepVariant這套工具看起來簡單實(shí)際規(guī)劃好并行粒度和中間文件清理跑大批量樣本的時(shí)候穩(wěn)定性和資源利用率都能提升一個(gè)檔次。如果你現(xiàn)在正打算把DeepVariant集成到自己的分析流程里可以先拿一個(gè)樣本把這個(gè)三階段流水線完整跑通再慢慢擴(kuò)展到集群級部署。等你的pipeline穩(wěn)定之后你會(huì)發(fā)現(xiàn)變異檢測這塊已經(jīng)從調(diào)參手藝人變成了流程工程師——后者省下來的時(shí)間才是最有價(jià)值的回報(bào)。