深度學習是一門很難評估代碼正確性的學科。隨機初始化、龐大的數據集和權重的有限可解釋性意味著，要找到模型為什么不能訓練的確切問題，大多數時候都需要反復試驗。在傳統的軟件開發中，自動化單元測試是確定代碼是否完成預期任務的面包和黃油。它幫助開發人員信任他們的代碼，并在引入更改時更加自信。一個破壞性的更改將會被單元測試檢測到。

[[338647]]

從GitHub上許多研究庫的情況來看，深度學習的實踐者們還不喜歡這種方法。從業者不知道他們的代碼是否正常工作，他們能接受嗎?通常，由于上述三個原因，學習系統的每個組件的預期行為并不容易定義。然而，我相信實踐者和研究人員應該重新考慮他們對單元測試的厭惡，因為它可以幫助研究過程更加順利。你只需要學習如何信任你的代碼。

顯然，我不是第一個，也不是最后一個談論用于深度學習的單元測試的人。如果你對這個話題感興趣，你可以看看這里：

• A Recipe for Training Neural Networks by Andrej Karpathy
• How to Unit Test Deep Learning by Sergios Karagiannakos

這篇文章的靈感來自于上面提到的，可能還有很多我現在想不起來的。為了在討論中增加一些內容，我們將重點關注如何編寫可重用的單元測試，這樣就可以“不去自己重復自己“。

我們的例子將測試用PyTorch編寫的系統的組件，該系統在MNIST 上訓練可變自動編碼器(VAE)。你可以在github.com/tilman151/unittest_dl上找到本文中的所有代碼。

什么是單元測試?

如果您熟悉單元測試，可以跳過此部分。對于其他人，我們將看到Python中的單元測試是什么樣子的。為了簡單起見，我們將使用內置的包unittest，而不是其他花哨的包。

一般來說，單元測試的目的是檢查代碼是否正確地運行。通常(我也為此感到內疚很長一段時間)，你會看到這樣的東西在一個文件的結尾：

 

if __name__ == 'main': 
    net = Network() 
    x = torch.randn(4, 1, 32, 32) 
    y = net(x) 
    print(y.shape) 

如果直接執行該文件，則代碼片段將構建一個網絡，執行前向傳遞并打印輸出的形狀。這樣，我們就可以看到向前傳播是否會拋出錯誤，以及輸出的形狀是否可信。如果將代碼分發到不同的文件中，則必須手動運行每個文件，并檢查打印到控制臺的內容。更糟糕的是，這個代碼片段有時會在運行后被刪除，當有變化時被重寫。

原則上，這已經是一個基本的單元測試。我們所要做的就是將它形式化一點，使它能夠輕松地自動運行。它看起來是這樣的：

 

import unittest 
 
class MyFirstTest(unittest.TestCase): 
    def test_shape(self): 
        net = Network() 
        x = torch.randn(4, 1, 32, 32) 
        y = net(x) 
        self.assertEqual(torch.Size((10,)), y.shape) 

unittest包的主要組件是類TestCase。單個單元測試是TestCase子類的成員函數。在我們的例子中，包將自動檢測類MyFirstTest并運行函數'test_shape。如果滿足assertEqual調用的條件，則測試成功。否則，或者如果它崩潰，測試將失敗。

我需要測試些什么?

現在我們已經了解了單元測試是如何工作的，下一個問題是我們應該測試什么。下面你可以看到我們的例子的代碼結構：

 

|- src 
   |- dataset.py 
   |- model.py 
   |- trainer.py 
   |- run.py 

我們將測試每個文件中的功能除了run.py，因為它只是我們程序的入口點。

Dataset

我們在例子中使用的數據集是torchvisionMNIST類。因此，我們可以假設像加載圖像和訓練/測試分割這樣的基本功能可以正常工作。然而，MNIST類為配置提供了充足的機會，因此我們應該測試是否正確配置了所有內容。dataset.py文件包含一個名為MyMNIST的類，它有兩個成員變量。成員train_data有torchvisionMNIST類的一個實例，該實例被配置為加載數據的訓練部分，而test_data 中的實例加載測試部分。兩種方法都將每幅圖像每邊填充2個像素，并將像素值歸一化在[- 1,1]之間。此外，train_data 對每個圖像應用隨機旋轉來增強數據。

數據的形狀

為了繼續使用上面的代碼片段，我們將首先測試數據集是否輸出了我們想要的形狀。圖像的填充意味著，它們現在的大小應該是32x32像素。我們的測試看起來是這樣的：

 

def test_shape(self): 
    dataset = MyMNIST() 
    sample, _ = dataset.train_data[0] 
    self.assertEqual(torch.Shape((1, 32, 32)), sample.shape) 

現在我們可以確定我們的padding是我們想要的。這可能看起來很瑣碎，你們中的一些人可能會認為我在測試這個方面很迂腐，但是我不知道我有多少次因為我搞不清楚填充函數是如何工作的而導致了形狀錯誤。像這樣的簡單測試編寫起來很快，并且可以為你以后省去許多麻煩。

數據的縮放

我們配置的下一件事是數據的縮放。在我們的例子中，這非常簡單。我們希望確保每個圖像的像素值在[- 1,1]之間。與之前的測試相反，我們將對數據集中的所有圖像進行測試。通過這種方式，我們可以確定我們關于如何縮放數據的假設對于整個數據集是有效的。

 

def test_scaling(self): 
    dataset = MyMNIST() 
    for sample, _ in dataset.train_data: 
        self.assertGreaterEqual(1, sample.max()) 
        self.assertLessEqual(-1, sample.min()) 
        self.assertTrue(torch.any(sample < 0)) 
        self.assertTrue(torch.any(sample > 0)) 

如你所見，我們不僅要測試每個圖像的最大值和最小值是否在范圍內。我們還通過斷言測試是否存在大于零和小于零的值，我們將值縮放到[0,1]。這個測試之所以有效，是因為我們可以假設MNIST中的每個圖像都覆蓋了整個范圍的值。對于更復雜的數據，比如自然圖像，我們需要一個更復雜的測試條件。如果你的縮放基于數據的統計信息，那么測試一下是否只使用訓練部分來計算這些統計信息也是一個好主意。

數據增強

增加訓練數據可以極大地幫助提高模型的性能，特別是在數據量有限的情況下。另一方面，我們不會增加我們的測試數據，因為我們想要保持我們的模型的評估確定性。這意味著，我們應該測試我們的訓練數據是否增加了，而我們的測試數據沒有。敏銳的讀者會在這一點上注意到一些重要的東西。到目前為止，我們的測試只涵蓋了訓練數據。這是需要強調的一點：

始終在訓練和測試數據上運行測試

僅僅因為你的代碼在數據的一個部分上工作，并不能保證在另一個部分上不存在未檢測到的bug。對于數據增強，我們甚至希望為每個部分斷言代碼的不同行為。

對于我們的增強問題，一個簡單的測試現在是加載一個樣本兩次，然后檢查兩個版本是否相等。簡單的解決方案是為我們的每一個部分寫一個測試函數：

 

def test_augmentation_active_train_data(self): 
    dataset = MyMNIST() 
    are_same = [] 
    for i in range(len(dataset.train_data)): 
        sample_1, _ = dataset.train_data[i] 
        sample_2, _ = dataset.train_data[i] 
        are_same.append(0 == torch.sum(sample_1 - sample_2)) 
 
    self.assertTrue(not all(are_same)) 
 
def test_augmentation_inactive_test_data(self): 
    dataset = MyMNIST() 
    are_same = [] 
    for i in range(len(dataset.test_data)): 
        sample_1, _ = dataset.test_data[i] 
        sample_2, _ = dataset.test_data[i] 
        are_same.append(0 == torch.sum(sample_1 - sample_2)) 
 
    self.assertTrue(all(are_same)) 

這些函數測試我們想要測試的內容，但是，正如你所看到的，它們幾乎就是重復的。這有兩個主要的缺點。首先，如果在測試中需要更改某些內容，我們必須記住在兩個函數中都要更改。其次，如果我們想添加另一個部分，例如一個驗證部分，我們將不得不第三次復制測試。要解決這個問題，我們應該將測試功能提取到一個單獨的函數中，然后由真正的測試函數調用兩次。重構后的測試看起來像這樣：

 

def test_augmentation(self): 
    dataset = MyMNIST() 
    self._check_augmentation(dataset.train_data, active=True) 
    self._check_augmentation(dataset.test_data, active=False) 
 
def _check_augmentation(self, data, active): 
    are_same = [] 
    for i in range(len(data)): 
        sample_1, _ = data[i] 
        sample_2, _ = data[i] 
        are_same.append(0 == torch.sum(sample_1 - sample_2)) 
 
    if active: 
        self.assertTrue(not all(are_same)) 
    else: 
        self.assertTrue(all(are_same)) 

_check_augmentation函數斷言給定的數據集是否進行了增強，并有效地刪除代碼中的重復。函數本身不會由unittest包自動運行，因為它不是以test_開頭的。因為我們的測試函數現在真的很短，我們把它們合并成一個組合函數。它們測試了增強是如何工作的這一單一的概念，因此應該屬于相同的測試函數。但是，通過這個組合，我們引入了另一個問題。如果測試失敗了，現在很難直接看到哪一個部分失敗了。這個包只告訴我們組合函數的名稱。進入subTest函數。TestCase類有一個成員函數subTest，它可以在一個測試函數中標記不同的測試組件。這樣，包就可以準確地告訴我們測試的哪一部分失敗了。最后的函數是這樣的：

 

def test_augmentation(self): 
    dataset = MyMNIST() 
    with self.subTest(split='train'): 
        self._check_augmentation(dataset.train_data, active=True) 
    with self.subTest(split='test'): 
        self._check_augmentation(dataset.test_data, active=False) 

現在我們有了一個無重復、精確定位、可重用的測試功能。我們在此所使用的核心原則可以應用到我們在前面幾節中編寫的所有其他單元測試中。你可以在附帶的存儲庫中看到結果測試。

數據的加載

數據集的最后一種類型的單元測試與我們的例子并不完全相關，因為我們使用的是內置數據集。無論如何我們都會把它包括進來，因為它涵蓋了我們學習系統的一個重要部分。通常，你將在dataloader類中使用數據集，該類處理批處理并可以并行化加載。因此，測試你的數據集在單進程和多進程模式下是否與dataloader一起工作是一個好主意。考慮到我們所學到的增強測試，測試函數如下所示：

 

def test_single_process_dataloader(self): 
    dataset = MyMNIST() 
    with self.subTest(split='train'): 
        self._check_dataloader(dataset.train_data, num_workers=0) 
    with self.subTest(split='test'): 
        self._check_dataloader(dataset.test_data, num_workers=0) 
 
def test_multi_process_dataloader(self): 
    dataset = MyMNIST() 
    with self.subTest(split='train'): 
        self._check_dataloader(dataset.train_data, num_workers=2) 
    with self.subTest(split='test'): 
        self._check_dataloader(dataset.test_data, num_workers=2) 
 
def _check_dataloader(self, data, num_workers): 
    loader = DataLoader(data, batch_size=4, num_workers=num_workers) 
    for _ in loader: 
        pass 

函數_check_dataloader不會對加載的數據進行任何測試。我們只是想檢查加載過程是否沒有拋出錯誤。理論上，ni 也可以檢查諸如正確的批大小或填充的序列數據的不同長度。因為我們為dataloader使用了最基本的配置，所以可以省略這些檢查。

同樣，這個測試可能看起來瑣碎而沒有必要，但是讓我給你一個例子，在這個簡單的檢查中節省了我的時間。這個項目需要從pandas的dataframes中加載序列數據，并從這些datafames上的滑動窗口中構造樣本。我們的數據集太大了，無法裝入內存，所以我們必須按需加載數據模型，并從中剪切出所請求的序列。為了提高加載速度，我們決定用一個LRU cache來緩存一些數據文件。它在我們早期的單進程實驗中如預期的那樣工作，因此我們決定將它包含在代碼庫中。結果是，這個緩存不能很好地用于多進程，但是我們的單元測試提前發現了這個問題。在使用多進程時，我們停用了緩存，避免了以后出現令人不快的意外。

最后要注意的

有些人可能已經在我們的單元測試中看到了另一個重復的模式。每個測試對訓練數據運行一次，對測試數據運行一次，產生相同的四行代碼：

 

with self.subTest(split='train'): 
    self._check_something(dataset.train_data) 
with self.subTest(split='test'): 
    self._check_dataloader(dataset.test_data) 

也完全有理由消除這種重復。不幸的是，這將涉及到創建一個高階函數，以函數_check_something作為參數。有時，例如對于增強測試，我們還需要向_check_something函數傳遞額外的參數。最后，所需的編程構造將引入更多的復雜性，并模糊要測試的概念。一般的規則是，為了可讀性和可重用性，讓你的測試代碼盡可能在需要的范圍內變復雜。

Model

模型可以說是學習系統的核心組件，通常需要是完全可配置的。這意味著，還有很多東西需要測試。幸運的是，PyTorch中用于神經網絡模型的API非常簡潔，大多數實踐者都非常嚴格地使用它。這使得為模型編寫可重用的單元測試相當容易。

我們的模型是一個簡單的VAE，由一個全連接的編碼器和解碼器組成。前向函數接受輸入圖像，對其進行編碼，執行重新參數化操作，然后將隱編碼解碼為圖像。雖然相對簡單，但這種變換可以演示幾個值得進行單元測試的方面。

模型的輸出形狀

我們在本文開頭看到的第一段代碼是幾乎每個人都要做的測試。我們也已經知道這個測試是如何寫成單元測試的。我們要做的唯一一件事就是添加要測試的正確形狀。對于一個自動編碼器，就簡單的判斷和輸入的形狀是否相同：

 

@torch.nograd() 
def test_shape(self): 
    net = model.MLPVAE(input_shape=(1, 32, 32), bottleneck_dim=16) 
    inputs = torch.randn(4, 1, 32, 32) 
    outputs = net(x) 
    self.assertEqual(inputs.shape, outputs.shape) 

同樣，這很簡單，但有助于找到一些最惱人的bug。例如，在將模型輸出從拉平的表示中reshape時忘記添加通道維度。

我們最后增加的測試是torch.nograd 。它告訴PyTorch這個函數不需要記錄梯度，并給我們一個小的加速。對于每個測試來說，它可能不是很多，但是你永遠不知道需要編寫多少。同樣，這是另一個可引用的單元測試智慧：

讓你的測試更快。否則，沒有人會想要運行它們。

單元測試應該在開發期間非常頻繁地運行。如果你的測試運行時間很長，那么你可以跳過它們。

模型的移動

在CPU上訓練深度神經網絡在大多數時候都非常慢。這就是為什么我們使用GPU來加速它。為此，我們所有的模型參數必須駐留在GPU上。因此，我們應該斷言我們的模型可以在設備(CPU和多個GPU)之間正確地移動。

我們可以用一個常見的錯誤來說明我們的例子VAE中的問題。這里你可以看到bottleneck函數，執行重新參數化的技巧：

 

def bottleneck(self, mu, log_sigma): 
    noise = torch.randn(mu.shape) 
    latent_code = log_sigma.exp() * noise + mu 
 
    return latent_code 

它取隱先驗的參數，從標準高斯分布中采樣一個噪聲張量，并使用參數對其進行變換。這在CPU上運行沒有問題，但當模型移動到GPU時失敗。問題是噪音張量是在CPU內存中創建的，因為它是默認的，并沒有移動到模型所在的設備上。一個簡單的錯誤和一個簡單的解決方案。我們用noise = torch.randn_like(mu)替換了這行有問題的代碼。這就產生了一個與張量mu相同形狀和在相同設備上的噪聲張量。

幫助我們盡早捕獲這些bug的測試：

 

@torch.no_grad() 
@unittest.skipUnless(torch.cuda.is_available(), 'No GPU was detected') 
def test_device_moving(self): 
    net = model.MLPVAE(input_shape=(1, 32, 32), bottleneck_dim=16) 
    net_on_gpu = net.to('cuda:0') 
    net_back_on_cpu = net_on_gpu.cpu() 
     
    inputs = torch.randn(4, 1, 32, 32) 
 
    torch.manual_seed(42) 
    outputs_cpu = net(inputs) 
    torch.manual_seed(42) 
    outputs_gpu = net_on_gpu(inputs.to('cuda:0')) 
    torch.manual_seed(42) 
    outputs_back_on_cpu = net_back_on_cpu(inputs) 
 
    self.assertAlmostEqual(0., torch.sum(outputs_cpu - outputs_gpu.cpu())) 
    self.assertAlmostEqual(0., torch.sum(outputs_cpu - outputs_back_on_cpu)) 

我們把網絡從一個CPU移動到另一個CPU，然后再移動回來，只是為了確保正確。現在我們有了網絡的三份拷貝(移動網絡復制了它們)，并使用相同的輸入張量向前傳遞。如果網絡被正確移動，前向傳遞應該在不拋出錯誤的情況下運行，并且每次產生相同的輸出。

為了運行這個測試，我們顯然需要一個GPU，但也許我們想在筆記本電腦上做一些快速測試。如果PyTorch沒有檢測到GPU，unittest.skipUnless 可以跳過測試。這樣可以避免將測試結果與失敗的測試混淆。

你還可以看到，我們在每次通過之前固定了torch的隨機種子。我們必須這樣做，因為VAEs是非確定性的，否則我們會得到不同的結果。這說明了深度學習代碼單元測試的另一個重要概念：

在測試中控制隨機性。

如果你不能確保你的模型能到邊界情況，你如何測試你的模型的一個罕見邊界條件?如何確保模型的輸出是確定性的?你如何知道一個失敗的測試是由于隨機的偶然還是由于你引入的bug ?通過手動設置深度學習框架的種子，可以消除函數中的隨機性。此外，還應該將CuDNN設置為確定性模式。這主要影響卷積，但無論如何是一個好主意。

注意確定正在使用的所有框架的種子。Numpy和內置的Python隨機數生成器有它們自己的種子，必須分別設置。有一個這樣的函數是很有用的：

 

def make_deterministic(seed=42): 
    # PyTorch 
    torch.manual_seed(seed) 
    if torch.cuda.is_available(): 
        torch.backends.cudnn.deterministic = True 
        torch.backends.cudnn.benchmark = False 
     
    # Numpy 
    np.random.seed(seed) 
     
    # Built-in Python 
    random.seed(seed) 

模型到采樣獨立性

在99。99%的情況下，你都想用隨機梯度下降的方式來訓練你的模型。你給你的模型一個minibatch的樣本，并計算他們的平均損失。批量處理訓練樣本假設你的模型可以處理每個樣本，也就是你可以獨立的把樣本喂給模型。換句話說，你的batch中的樣本在你的模型處理時不會相互影響。這個假設是很脆弱的，如果在一個錯誤的張量維度上進行錯誤的reshape或aggregation，就會打破這個假設。

下面的測試通過執行與輸入相關的前向和后向傳遞來檢查樣本的獨立性。在對這個batch做平均損失之前，我們把損失乘以零。如果我們的模型保持樣本獨立性，這將導致一個零梯度。唯一的事情，我們必須斷言，如果只有masked的樣本梯度是零：

 

def test_batch_independence(self): 
    inputs = torch.randn(4, 1, 32, 32) 
    inputs.requires_grad = True 
    net = model.MLPVAE(input_shape=(1, 32, 32), bottleneck_dim=16) 
 
    # Compute forward pass in eval mode to deactivate batch norm 
    net.eval() 
    outputs = net(inputs) 
    net.train() 
 
    # Mask loss for certain samples in batch 
    batch_size = inputs[0].shape[0] 
    mask_idx = torch.randint(0, batch_size, ()) 
    mask = torch.ones_like(outputs) 
    mask[mask_idx] = 0 
    outputs = outputs * mask 
 
    # Compute backward pass 
    loss = outputs.mean() 
    loss.backward() 
 
    # Check if gradient exists and is zero for masked samples 
    for i, grad in enumerate(inputs.grad): 
        if i == mask_idx: 
            self.assertTrue(torch.all(grad == 0).item()) 
        else: 
            self.assertTrue(not torch.all(grad == 0)) 

如果你準確地閱讀了代碼片段，你會注意到我們將模型設置為evaluation模式。這是因為batch normalization違反了我們上面的假設。進程均值和標準差的處理交叉污染了我們batch中的樣本，所以我們通過evaluation模式停止了對樣本的更新。我們可以這樣做，因為我們的模型在訓練和評估模式中表現相同。如果你的模型不是這樣的，你將不得不找到另一種方法來禁用它進行測試。一個選項是用instance normalization臨時替換它。

上面的測試函數非常通用，可以按原樣復制。例外情況是，如果你的模型接受多個輸入。處理這個問題的附加代碼是必要的。

模型的參數更新

下一個測試也與梯度有關。當你的網絡架構變得更加復雜時，比如初始化，很容易構建死子圖。死子圖是網絡中包含可學習參數的一部分，前向傳遞、后向傳遞或兩者都不使用。這就像在構造函數中構建一個網絡層，然后忘記在forward函數中應用它一樣簡單。

找到這些死子圖可以通過運行優化步驟并檢查梯度你的網絡參數：

 

def test_all_parameters_updated(self): 
    net = model.MLPVAE(input_shape=(1, 32, 32), bottleneck_dim=16) 
    optim = torch.optim.SGD(net.parameters(), lr=0.1) 
 
    outputs = net(torch.randn(4, 1, 32, 32)) 
    loss = outputs.mean() 
    loss.backward() 
    optim.step() 
 
    for param_name, param in self.net.named_parameters(): 
        if param.requires_grad: 
            with self.subTest(name=param_name): 
                self.assertIsNotNone(param.grad) 
                self.assertNotEqual(0., torch.sum(param.grad ** 2)) 

參數函數返回的模型的所有參數在優化步驟后都應該有一個梯度張量。此外，對于我們所使用的損失，它不應該是零。測試假設模型中的所有參數都需要梯度。即使是那些不應該被更新的參數也會首先檢查requires_grad標志。如果任何參數在測試中失敗，子測試的名稱將提示你在哪里查找。

提高重用性

現在我們已經寫出了模型的所有測試，我們可以將它們作為一個整體進行分析。我們將注意到這些測試有兩個共同點。所有測試都從創建模型和定義示例輸入批處理開始。與以往一樣，這種冗余級別有可能導致拼寫錯誤和不一致。此外，你不希望在更改模型的構造函數時分別更新每個測試。

幸運的是，unittest為我們提供了一個簡單的解決方案，即setUp函數。這個函數在執行TestCase中的每個測試函數之前被調用，通常為空。通過在setUp中將模型和輸入定義為TestCase的成員變量，我們可以在一個地方初始化測試的組件。

 

class TestVAE(unittest.TestCase): 
    def setUp(self): 
        self.net = model.MLPVAE(input_shape=(1, 32, 32), bottleneck_dim=16) 
        self.test_input = torch.random(4, 1, 32, 32) 
 
    ... # Test functions 

現在我們用各自的成員變量替換出現的net和inputs，這樣就完成了。如果你想更進一步，對所有測試使用相同的模型實例，您可以使用setUpClass。這個函數在構造TestCase時被調用一次。如果構建速度很慢，并且你不想多次進行構建，那么這是非常有用的。

在這一點上，我們有一個整潔的系統來測試我們的VAE模型。我們可以輕松地添加測試，并確保每次都測試模型的相同版本。但是如果你想引入一種新的卷積層，會發生什么呢?它將在相同的數據上運行，也應該具有相同的行為，因此將應用相同的測試。

僅僅復制整個TestCase 顯然不是首選的解決方案，但是通過使用setUp，我們已經在正確的軌道上了。我們將所有測試函數轉移到一個基類中，而將setUp保留為一個抽象函數。

 

class AbstractTestVAE(unittest.TestCase): 
    def setUp(self): 
        raise NotImplementedError 
 
    ... # Test functions 

你的IDE會提示類沒有成員變量net 和test_inputs，但是Python并不關心。只要子類添加了它們，它就可以工作。對于我們想要測試的每個模型，我們創建這個抽象類的一個子類，并在其中實現setUp。為多個模型或同一個模型的多個配置創建TestCases 就像：

 

class TestCNNVAE(AbstractTestVAE): 
    def setUp(self): 
        self.test_inputs = torch.randn(4, 1, 32, 32) 
        self.net = model.CNNVAE(input_shape=(1, 32, 32), bottleneck_dim=16) 
 
class TestMLPVAE(AbstractTestVAE): 
    def setUp(self): 
        self.test_inputs = torch.randn(4, 1, 32, 32) 
        self.net = model.MLPVAE(input_shape=(1, 32, 32), bottleneck_dim=16) 

只剩下一個問題了。unittest包發現并運行unittest.TestCase的所有子元素。因為這包括不能實例化的抽象基類，所以我們總是會有一個失敗的測試。

解決方案是由一個流行的設計模式提出的。通過刪除TestCase作為AbstractTestVAE的父類，它就不再被發現了。相反，我們讓我們的具體測試有兩個父類， TestCase和AbstractTestVAE。抽象類和具體類之間的關系不再是父類和子類之間的關系。相反，具體類使用抽象類提供的共享功能。這個模式稱為MixIn。

 

class AbstractTestVAE: 
    ... 
 
class TestCNNVAE(unittest.TestCase, AbstractTestVAE): 
    ... 
 
class TestMLPVAE(unittest.TestCase, AbstractTestVAE): 
    ... 

父類的順序很重要，因為方法查找是從左到右進行的。這意味著TestCase將覆蓋AbstractTestVAE的共享方法。在我們的例子中，這不是一個問題，但無論如何知道都是好的。

Trainer

我們的學習系統的最后一部分是trainer類。它將你所有的組件(數據集、優化器和模型)放在一起，并使用它們來訓練模型。此外，它還實現了一個評估函數，輸出測試數據的平均損失。在訓練時，所有的損失和指標都被寫入一個TensorBoard event文件中以便可視化。

在這一部分中，編寫可重用測試是最困難的，因為它允許最大程度的自由實現。有些人只在腳本文件中使用簡單的代碼進行訓練，有些人將其封裝在函數中，還有一些人試圖保持更面向對象的風格。我不會判斷你喜歡哪種方式。我唯一要說的是，在我的經驗中，整潔封裝的trainer類使單元測試變得最舒適。

然而，我們會發現我們之前學過的一些原則在這里也適用。

trainer的損失

大多數時候，你只需要從torch上選擇一個預先實現的損失函數就可以了。但話說回來，你所選擇的損失函數可能無法實現。這種情況可能是由于實現相對簡單，函數太小眾或者太新。無論如何，如果你自己實現了它，你也應該測試它。

我們的例子使用Kulback-Leibler (KL)散度作為整體損失函數的一部分，這在PyTorch中是不存在的(現在的版本里有了)。我們的實現是這樣的：

 

def _kl_divergence(log_sigma, mu): 
    return 0.5 * torch.sum((2 * log_sigma).exp() + mu ** 2 - 1 - 2 * log_sigma) 

函數取多變量高斯分布的標準偏差和平均值的對數，并計算在封閉形式中的標準高斯分布的KL散度。

檢查這種損失的一種方法是手工計算，然后硬編碼以便比較。更好的方法是在另一個包中找到一個參考實現，并根據它的輸出檢查代碼。幸運的是，scipy包有一個離散KL散度的實現，我們可以使用：

 

@torch.no_grad() 
def test_kl_divergence(self): 
    mu = np.random.randn(10) * 0.25  # means around 0. 
    sigma = np.random.randn(10) * 0.1 + 1.  # stds around 1. 
    standard_normal_samples = np.random.randn(100000, 10) 
    transformed_normal_sample = standard_normal_samples * sigma + mu 
 
    bins = 1000 
    bin_range = [-2, 2] 
    expected_kl_div = 0 
    for i in range(10): 
        standard_normal_dist, _ = np.histogram(standard_normal_samples[:, i], bins, bin_range) 
        transformed_normal_dist, _ = np.histogram(transformed_normal_sample[:, i], bins, bin_range) 
        expected_kl_div += scipy.stats.entropy(transformed_normal_dist, standard_normal_dist) 
 
    actual_kl_div = self.vae_trainer._kl_divergence(torch.tensor(sigma).log(), torch.tensor(mu)) 
 
    self.assertAlmostEqual(expected_kl_div, actual_kl_div.numpy(), delta=0.05) 

我們首先從標準高斯函數和一個不同均值和標準差的高斯函數中抽取一個足夠大的樣本。然后我們用np.histogram函數，得到基本pdf的離散逼近。有了這些，我們就可以用scipy.stats.entropy得到一個KL散度來比較。我們使用一個相對較大的delta來進行比較，因為scipy.stats.entropy只是一個近似值。

你可能已經注意到，我們沒有創建Trainer對象，而是使用TestCase的成員。我們在這里使用了與模型測試相同的技巧，并在setUp函數中創建了它。我們還固定了PyTorch和NumPy的種子。因為我們這里不需要任何梯度，所以我們用@torch.no_grad來裝飾函數。

trainer的日志記錄

我們使用TensorBoard來記錄我們的訓練過程的損失和度量。為此，我們希望確保按預期寫入所有日志。一種方法是在訓練后打開event文件，查找正確的event。同樣，這也是一個有效的選項，但我們將以另一種方式來看看unittest包的一個有趣功能：mock。

mock允許你用一個監視其自身是如何調用的函數來打包一個函數或對象。我們將替換summary writer的add_scalar 函數，并確保以這種方式記錄我們關心的所有損失和指標。

 

def test_logging(self): 
    with mock.patch.object(self.vae_trainer.summary, 'add_scalar') as add_scalar_mock: 
        self.vae_trainer.train(1) 
 
    expected_calls = [mock.call('train/recon_loss', mock.ANY, 0), 
                      mock.call('train/kl_div_loss', mock.ANY, 0), 
                      mock.call('train/loss', mock.ANY, 0), 
                      mock.call('test/loss', mock.ANY, 0)] 
    add_scalar_mock.assert_has_calls(expected_calls) 

assert_has_calls 函數匹配預期調用列表和實際記錄的調用。mock.ANY 表示我們不關心記錄的標量的值，因為無論如何我們都不知道它。

因為我們不需要對整個數據集執行完一個epoch，所以我們在setUp 中將訓練數據配置為只有一個batch。這樣，我們可以顯著地加快我們的測試速度。

trainer的擬合

最后一個問題也是最難回答的。我的訓練最終會收斂嗎?要確切地回答這個問題，我們需要用我們所有的數據進行一次全面的訓練并對其打分。

由于這非常耗時，我們將使用一種更快的方法。我們將看看我們的訓練是否能使模型對單個batch的數據進行過擬合。測試函數相當簡單：

 

def test_overfit_on_one_batch(self): 
    self.vae_trainer.train(500) 
    self.assertGreaterEqual(30, self.vae_trainer.eval()) 

如前一節所述，setUp函數創建一個只包含一個batch的數據集的trainer。此外，我們也使用訓練數據作為測試數據。通過這種方式，我們可以從 eval函數中獲得訓練batch的損失，并將其與我們預期的損失進行比較。

對于一個分類問題，當我們完全過擬合時，我們期望損失為零。“VAE”的問題是，它是一個非確定性的生成模型，零損失是不現實的。這就是為什么我們預期的損失是30，這等于每像素的誤差為0.04。

這是迄今為止運行時間最長的測試，它可以運行500 epochs。最后，在我的筆記本電腦上用1.5分鐘左右就可以了，這仍然是合理的。為了在不降低對沒有GPU的機器的支持的情況下進一步加速，我們可以簡單地在setUp中添加這一行：

 

device = 'cuda:0' if torch.cuda.is_available() else 'cpu' 

這樣一來，如果我們有GPU，我們就可以利用它，如果沒有，就利用CPU進行訓練。

最后要注意的

在我們進行日志記錄時，你可能會注意到，針對trainer的單元測試往往會使你的文件夾充滿event文件。為了避免這種情況，我們使用tempfile 包為trainer創建一個臨時日志目錄。測試結束后，我們只需要再次刪除它和它的內容。為此，我們使用了孿生函數setUp,和tearDown。在每個測試函數后調用此函數，清理過程簡單如下：

 

def tearDown(self): 
    shutil.rmtree(self.log_dir) 

總結

我們看完了這篇文章。讓我們評估一下我們從整個磨難中得到了什么。

我們為我們的小例子編寫的測試套件包含58個單元測試，整個運行大約需要3.5分鐘。對于這58個測試，我們只編寫了20個函數。所有測試都可以確定地、獨立地運行。如果有GPU，我們可以運行額外的測試。大多數測試，例如數據集和模型測試，可以在其他項目中輕松重用。我們可以通過使用：

• 子測試為我們的數據集的多種配置運行一個測試
• setUp和tearDown函數一致地初始化和清理我們的測試
• 抽象測試類來測試VAE的不同實現
• torch.no_grad裝飾器在可能的情況下禁用梯度計算
• mock模塊檢查函數是否被正確調用

最后，我希望我能夠說服至少有人在他們的深度學習項目中使用單元測試。本文的配套git倉庫可以作為起點。