Keras 和 PyTorch 都是對(duì)初學(xué)者非常友好的深度學(xué)習(xí)框架，兩者各有優(yōu)勢(shì)，很多研究者和開(kāi)發(fā)者在選擇框架時(shí)可能會(huì)舉棋不定?；谶@種情況，grid.ai CEO、紐約大學(xué)博士 William Falcon 創(chuàng)建了 PyTorch Lightning，為 PyTorch 披上了一件 Keras 的外衣。

[[347903]]

Lightning 是 PyTorch 非常輕量級(jí)的包裝，研究者只需要編寫最核心的訓(xùn)練和驗(yàn)證邏輯，其它過(guò)程都會(huì)自動(dòng)完成。因此這就有點(diǎn)類似 Keras 那種高級(jí)包裝，它隱藏了絕大多數(shù)細(xì)節(jié)，只保留了最通俗易懂的接口。Lightning 能確保自動(dòng)完成部分的正確性，對(duì)于核心訓(xùn)練邏輯的提煉非常有優(yōu)勢(shì)。

今日，PyTorch Lightning 在推特宣布，1.0.0 版本現(xiàn)在可用了，并發(fā)布新的博客文章詳細(xì)描述了 PyTorch Lightning 的運(yùn)行原理和新的 API。William Falcon 表示自己非常期待有一天，當(dāng)用戶查看 GitHub 上的復(fù)雜項(xiàng)目時(shí)，深度學(xué)習(xí)代碼不再那么令人望而生畏。

特斯拉 AI 負(fù)責(zé)人 Andrej Karpathy 也評(píng)論稱：「這看起來(lái)很棒，也很有前途。PyTorch Lightning 倡導(dǎo)對(duì)深度學(xué)習(xí)代碼進(jìn)行重構(gòu)，將『工程(硬件)』與『科學(xué)(代碼)』分割開(kāi)，然后將前者委托給框架?！?/p>

過(guò)去幾個(gè)月里，PyTorch Lightning 團(tuán)隊(duì)一直在微調(diào) API、完善文檔和記錄教程，最終使得 V1.0.0 順利面世。在接下來(lái)的博客文章中，該團(tuán)隊(duì)對(duì) PyTorch Lightning 進(jìn)行了詳盡的解讀。

• 博客地址：https://medium.com/pytorch/pytorch-lightning-1-0-from-0-600k-80fc65e2fab0
• GitHub 地址：https://github.com/PyTorchLightning/pytorch-lightning

PyTorch Lightning 的運(yùn)行原理和目標(biāo)

人工智能的發(fā)展速度比單一框架發(fā)展要快得多。深度學(xué)習(xí)領(lǐng)域在不斷發(fā)展，主要體現(xiàn)在復(fù)雜度與規(guī)模性兩方面。Lightning 提供了一種為復(fù)雜模型交互設(shè)計(jì)的用戶體驗(yàn)，同時(shí)抽象化了工程中許多零散的細(xì)節(jié)，如多 GPU 和多 TPU 訓(xùn)練、提前停止、日志記錄等…

像 PyTorch 這樣的框架出現(xiàn)的時(shí)間，正是人工智能主要關(guān)注網(wǎng)絡(luò)架構(gòu)的階段。

這些框架提供所有的部件來(lái)組合極其復(fù)雜的模型，在研究和生產(chǎn)方面做得非常出色。但是，一旦模型開(kāi)始交互，像 GAN， BERT 或者是自動(dòng)編碼器，范式就被打破，很快就失去了極好的靈活性，很難按照項(xiàng)目規(guī)模進(jìn)行維護(hù)。

與之前的框架不同，PyTorch Lightning 用來(lái)封裝一系列相互作用的模型，即深度學(xué)習(xí)系統(tǒng)。Lightning 是為當(dāng)今世界更復(fù)雜的研究以及生產(chǎn)案例而建立的，在這種情況下，許多模型使用復(fù)雜的規(guī)則進(jìn)行交互。

自動(dòng)編碼系統(tǒng)

PyTorch Lightning 的第二個(gè)關(guān)鍵原理是硬件和科學(xué)代碼分開(kāi)。Lightning 的發(fā)展可以大規(guī)模地利用大量計(jì)算，而不會(huì)向用戶呈現(xiàn)任何抽象概念。通過(guò)這種分離，你可以獲得以前不可能實(shí)現(xiàn)的新功能，比如，無(wú)需更改代碼就可以在筆記本電腦上使用 CPU 調(diào)試 512 GPU。

最后，Lightning 希望成為一個(gè)社區(qū)驅(qū)動(dòng)的框架。

構(gòu)建良好的深度學(xué)習(xí)模型需要大量的專業(yè)知識(shí)和小技巧，才能使系統(tǒng)正常工作。在世界各地，數(shù)以百計(jì)的工程師和博士不斷地實(shí)現(xiàn)同樣的代碼。現(xiàn)在，Lightning 的貢獻(xiàn)者社區(qū)不斷壯大，有超過(guò) 300 名最具天賦的深度學(xué)習(xí)人士，他們選擇分配相同的能量并進(jìn)行完全相同的優(yōu)化，但卻有成千上萬(wàn)的人從他們的努力中受益。

[[347904]]

PyTorch Lightning 1.0.0 有哪些新功能

Lightning 1.0.0 標(biāo)志著一個(gè)穩(wěn)定的最終 API。這對(duì)使用 Lightning 的研究者來(lái)說(shuō)是一件好事，因?yàn)樗麄兊拇a不會(huì)輕易被破壞或改變。

1. 研究與生產(chǎn)

Lightning 的核心優(yōu)勢(shì)是：使得最先進(jìn)的人工智能研究能夠大規(guī)模進(jìn)行。這是一個(gè)為專業(yè)研究人員設(shè)計(jì)的框架，可以在最大的計(jì)算資源上嘗試最難的想法，而不會(huì)失去任何靈活性。

Lightning 1.0.0 使大規(guī)模的部署模型變得簡(jiǎn)單。代碼可以輕松導(dǎo)出。

這意味著數(shù)據(jù)科學(xué)家、研究人員等團(tuán)隊(duì)現(xiàn)在就可以成為生產(chǎn)模型的人，而不需要龐大的機(jī)器學(xué)習(xí)工程師團(tuán)隊(duì)。

Lightning 旨在提供一種幫助研究者大幅縮短生產(chǎn)時(shí)間的方法，同時(shí)又不喪失任何研究所需的靈活性

Grid AI 是用于在云上進(jìn)行大規(guī)模訓(xùn)練模型的本機(jī)平臺(tái)。該平臺(tái)允許構(gòu)建深度學(xué)習(xí)模型的研究者在大規(guī)模計(jì)算上進(jìn)行迭代，然后將模型部署到可擴(kuò)展環(huán)境中，該環(huán)境能夠處理深度學(xué)習(xí)系統(tǒng)的最大流量。

2. 度量指標(biāo)

pytorch_lightning.metrics 是一種 Metrics API，旨在在 PyTorch 和 PyTorch Lightning 中輕松地進(jìn)行度量指標(biāo)的開(kāi)發(fā)和使用。更新后的 API 提供了一種內(nèi)置方法，可針對(duì)每個(gè)步驟跨多個(gè) GPU(進(jìn)程)計(jì)算指標(biāo)，同時(shí)存儲(chǔ)統(tǒng)計(jì)信息。這可以讓用戶在一個(gè)階段結(jié)束時(shí)計(jì)算指標(biāo)，而無(wú)需擔(dān)心任何與分布式后端相關(guān)的復(fù)雜度。

class LitModel(pl.LightningModule): 
    def __init__(self): 
        ... 
        self.train_acc = pl.metrics.Accuracy() 
        self.valid_acc = pl.metrics.Accuracy() 
 
    def training_step(self, batch, batch_idx): 
        logits = self(x) 
        ... 
        self.train_acc(logits, y) 
        # log step metric 
        self.log('train_acc_step', self.train_acc) 
 
    def validation_step(self, batch, batch_idx): 
        logits = self(x) 
        ... 
        self.valid_acc(logits, y) 
        # logs epoch metrics 
        self.log('valid_acc', self.valid_acc) 

要實(shí)現(xiàn)自定義指標(biāo)，只需將 Metric 基類子類化，并實(shí)現(xiàn)__init__()、update() 和 compute() 方法。用戶需要做的就是正確調(diào)用 add_state()，以用 DDP 實(shí)現(xiàn)自定義指標(biāo)。對(duì)使用 add_state() 添加的度量指標(biāo)狀態(tài)變量要調(diào)用 reset()。

from pytorch_lightning.metrics import Metric 
 
class MyAccuracy(Metric): 
 
    def __init__(self, dist_sync_on_step=False): 
        super().__init__(dist_sync_on_stepdist_sync_on_step=dist_sync_on_step) 
        self.add_state("correct", default=torch.tensor(0), dist_reduce_fx="sum") 
        self.add_state("total", default=torch.tensor(0), dist_reduce_fx="sum") 
         
    def update(self, preds: torch.Tensor, target: torch.Tensor): 
        preds, target = self._input_format(preds, target) 
        assert preds.shape == target.shape 
        self.correct += torch.sum(preds == target) 
        self.total += target.numel() 
  
    def compute(self): 
        return self.correct.float() / self.total 

3. 手動(dòng)優(yōu)化 VS 自動(dòng)優(yōu)化

使用 Lightning，用戶不需要擔(dān)心何時(shí)啟用 / 停用 grad，只要從 training_step 中返回帶有附加圖的損失即可進(jìn)行反向傳播或更新優(yōu)化器，Lightning 將會(huì)自動(dòng)進(jìn)行優(yōu)化。

def training_step(self, batch, batch_idx): 
    loss = self.encoder(batch[0]) 
    return loss 

但是，對(duì)于某些研究，如 GAN、強(qiáng)化學(xué)習(xí)或者是帶有多個(gè)優(yōu)化器或內(nèi)部循環(huán)的某些研究，用戶可以關(guān)閉自動(dòng)優(yōu)化，并完全由自己控制訓(xùn)練循環(huán)。

首先，關(guān)閉自動(dòng)優(yōu)化：

trainer *=* Trainer(automatic_optimization*=False*) 

現(xiàn)在訓(xùn)練循環(huán)已經(jīng)由用戶自己掌握。

def training_step(self, batch, batch_idx, opt_idx): 
    (opt_a, opt_b, opt_c) = self.optimizers() 
    loss_a = self.generator(batch[0]) 
    # use this instead of loss.backward so we can automate half 
    # precision, etc... 
    self.manual_backward(loss_a, opt_a, retain_graph=True) 
    self.manual_backward(loss_a, opt_a) 
    opt_a.step() 
    opt_a.zero_grad() 
    loss_b = self.discriminator(batch[0]) 
    self.manual_backward(loss_b, opt_b) 
    ... 

4. Logging

Lightning 讓帶有 logger 的集成變得非常簡(jiǎn)單，只要在 LightningModule 中調(diào)用 log() 方法，系統(tǒng)就會(huì)將已記錄的數(shù)量發(fā)送到用戶選擇的 logger 上。默認(rèn)使用 Tensorboard，但是你也可以選擇其他支持的 logger。

def training_step(self, batch, batch_idx): 
  self.log('my_metric', x) 

根據(jù)調(diào)用. log() 的位置，Lightning 會(huì)自動(dòng)確定何時(shí)記錄(在每一步和每個(gè)階段)，但是用戶也可以通過(guò)手動(dòng)使用 on_step 和 on_epoch 來(lái)重寫默認(rèn)行為。

但是 on_epoch=True 將會(huì)在整個(gè)訓(xùn)練階段累積記錄值。

def training_step(self, batch, batch_idx): 
  self.log('my_loss', loss, on_step=True, on_epoch=True, prog_bar=True, logger=True) 

5. 數(shù)據(jù)流

Lightning 棄用了 EvaResult 和 TrainResult，這有利于簡(jiǎn)化數(shù)據(jù)流并將訓(xùn)練和驗(yàn)證循環(huán)(loop)中的日志記錄與數(shù)據(jù)解耦。

每個(gè)循環(huán)(訓(xùn)練、驗(yàn)證和測(cè)試)具有三個(gè)可以實(shí)現(xiàn)的 hook，具體如下：

• x_step
• x_step_end
• x_epoch_end

為了演示數(shù)據(jù)流的運(yùn)行方式，以下代碼實(shí)現(xiàn)中使用到了訓(xùn)練循環(huán)(即 x=training)：

outs *=* [] 
*for* batch *in* data: 
  out *=* training_step(batch) 
  outs*.*append(out)training_epoch_end(outs) 

在 training_step 中返回的任意內(nèi)容可以作為 training_epoch_end 的輸入。

def training_step(self, batch, batch_idx): 
  prediction = … 
  return {'loss': loss, 'preds': prediction} 
 
def training_epoch_end(self, training_step_outputs): 
  for out in training_step_outputs: 
    prediction = out['preds'] 
  # do something with these 

驗(yàn)證和測(cè)試循環(huán)的代碼實(shí)現(xiàn)也是同樣的步驟。如果想要使用 DP 或者 DDP2 分布式模式(即在 GPU 上分割 batch)，則使用 x_step_end 進(jìn)行手動(dòng)聚合(或者不實(shí)現(xiàn)，令 lightning 進(jìn)行自動(dòng)聚合)。

6. 檢查點(diǎn)

現(xiàn)在，Lightning 可以通過(guò)用戶最后訓(xùn)練 epoch 的狀態(tài)，在當(dāng)前工作目錄中自動(dòng)保存檢查點(diǎn)。這保證用戶可以在訓(xùn)練中斷的情況下重新開(kāi)始。

此外，用戶可以自定義檢查點(diǎn)行為，以監(jiān)控任意數(shù)量的訓(xùn)練或驗(yàn)證步驟。例如，如果用戶想要基于自己的驗(yàn)證損失來(lái)更新檢查點(diǎn)，則可以按照以下步驟完成：

• 計(jì)算想要監(jiān)控的任意度量或其他數(shù)量，如驗(yàn)證損失;
• 通過(guò) log() 方法記錄下數(shù)量以及 val_loss 等鍵(key);
• 初始化 ModelCheckpoint 回調(diào)函數(shù)，將 monitor 設(shè)置為數(shù)量的 key;
• 將回調(diào)函數(shù) checkpoint_callback 返回訓(xùn)練器 flag。

具體代碼過(guò)程如下：

from pytorch_lightning.callbacks import ModelCheckpoint 
 
class LitAutoEncoder(pl.LightningModule): 
    def validation_step(self, batch, batch_idx): 
        x, y = batch 
        y_hat = self.backbone(x) 
 
        # 1. calculate loss 
        loss = F.cross_entropy(y_hat, y) 
 
        # 2. log `val_loss` 
        self.log('val_loss', loss) 
 
# 3. Init ModelCheckpoint callback, monitoring 'val_loss' 
checkpoint_callback = ModelCheckpoint(monitor='val_loss') 
 
# 4. Pass your callback to checkpoint_callback trainer flag 
trainer = Trainer(checkpoint_callbackcheckpoint_callback=checkpoint_callback) 

所有 API 的變化，包括 bug 修復(fù)等可以參考 GitHub 項(xiàng)目。

【本文是51CTO專欄機(jī)構(gòu)“機(jī)器之心”的原創(chuàng)譯文，微信公眾號(hào)“機(jī)器之心( id: almosthuman2014)”】 

戳這里，看該作者更多好文