Pytorch：常用工具模塊

數據處理

在解決深度學習問題的過程中，往往需要花費大量的精力去處理數據，包括圖像、文本、語音或其它二進制數據等。數據的處理對訓練神經網絡來說十分重要，良好的數據處理不僅會加速模型訓練，更會提高模型效果。考慮到這點，PyTorch提供了幾個高效便捷的工具，以便使用者進行數據處理或增強等操作，同時可通過並行化加速數據加載。

數據加載

在PyTorch中，數據加載可通過自定義的數據集對象。數據集對象被抽象為Dataset類，實現自定義的數據集需要繼承Dataset，並實現兩個Python魔法方法：

• __getitem__：返回一條數據，或一個樣本。obj[index]等價於obj.__getitem__(index)
• __len__：返回樣本的數量。len(obj)等價於obj.__len__()

import torch as t
from torch.utils import data

import os
from PIL import  Image
import numpy as np

class DogCat(data.Dataset):
    def __init__(self, root):
        imgs = os.listdir(root)
        # 所有圖片的絕對路徑
        # 這里不實際加載圖片，只是指定路徑，當調用__getitem__時才會真正讀圖片
        self.imgs = [os.path.join(root, img) for img in imgs]
        
    def __getitem__(self, index):
        img_path = self.imgs[index]
        # dog->1， cat->0
        label = 1 if 'dog' in img_path.split('/')[-1] else 0
        pil_img = Image.open(img_path)
        array = np.asarray(pil_img)
        data = t.from_numpy(array)
        return data, label
    
    def __len__(self):
        return len(self.imgs)

通過上面的代碼，我們學習了如何自定義自己的數據集，並可以依次獲取。但這里返回的數據不適合實際使用，因其具有如下兩方面問題：

• 返回樣本的形狀不一，因每張圖片的大小不一樣，這對於需要取batch訓練的神經網絡來說很不友好
• 返回樣本的數值較大，未歸一化至[-1, 1]

針對上述問題，PyTorch提供了torchvision1。它是一個視覺工具包，提供了很多視覺圖像處理的工具，其中transforms模塊提供了對PIL Image對象和Tensor對象的常用操作

對PIL Image的操作包括：

• Scale：調整圖片尺寸，長寬比保持不變
• CenterCrop、RandomCrop、RandomResizedCrop： 裁剪圖片
• Pad：填充
• ToTensor：將PIL Image對象轉成Tensor，會自動將[0, 255]歸一化至[0, 1]

對Tensor的操作包括：

• Normalize：標准化，即減均值，除以標准差
• ToPILImage：將Tensor轉為PIL Image對象

如果要對圖片進行多個操作，可通過Compose函數將這些操作拼接起來，類似於nn.Sequential。注意，這些操作定義后是以函數的形式存在，真正使用時需調用它的__call__方法，這點類似於nn.Module。例如要將圖片調整為224×224224×224，首先應構建這個操作trans = Resize((224, 224))，然后調用trans(img)。下面我們就用transforms的這些操作來優化上面實現的dataset。

import os
from PIL import  Image
import numpy as np
from torchvision import transforms as T

transform = T.Compose([
    T.Resize(224), # 縮放圖片(Image)，保持長寬比不變，最短邊為224像素
    T.CenterCrop(224), # 從圖片中間切出224*224的圖片
    T.ToTensor(), # 將圖片(Image)轉成Tensor，歸一化至[0, 1]
    T.Normalize(mean=[.5, .5, .5], std=[.5, .5, .5]) # 標准化至[-1, 1]，規定均值和標准差
])

class DogCat(data.Dataset):
    def __init__(self, root, transforms=None):
        imgs = os.listdir(root)
        self.imgs = [os.path.join(root, img) for img in imgs]
        self.transforms=transforms
        
    def __getitem__(self, index):
        img_path = self.imgs[index]
        label = 0 if 'dog' in img_path.split('/')[-1] else 1
        data = Image.open(img_path)
        if self.transforms:
            data = self.transforms(data)
        return data, label
    
    def __len__(self):
        return len(self.imgs)

dataset = DogCat('./data/dogcat/', transforms=transform)
img, label = dataset[0]
for img, label in dataset:
    print(img.size(), label)

torchvision已經預先實現了常用的Dataset，包括前面使用過的CIFAR-10，以及ImageNet、COCO、MNIST、LSUN等數據集，可通過諸如torchvision.datasets.CIFAR10來調用，具體使用方法請參看官方文檔1。在這里介紹一個會經常使用到的Dataset——ImageFolder，它的實現和上述的DogCat很相似。ImageFolder假設所有的文件按文件夾保存，每個文件夾下存儲同一個類別的圖片，文件夾名為類名，其構造函數如下：

ImageFolder(root, transform=None, target_transform=None, loader=default_loader)

它主要有四個參數：

• root：在root指定的路徑下尋找圖片
• transform：對PIL Image進行的轉換操作，transform的輸入是使用loader讀取圖片的返回對象
• target_transform：對label的轉換
• loader：給定路徑后如何讀取圖片，默認讀取為RGB格式的PIL Image對象

label是按照文件夾名順序排序后存成字典，即{類名:類序號(從0開始)}，一般來說最好直接將文件夾命名為從0開始的數字，這樣會和ImageFolder實際的label一致，如果不是這種命名規范，建議看看self.class_to_idx屬性以了解label和文件夾名的映射關系。

Dataset只負責數據的抽象，一次調用__getitem__只返回一個樣本。前面提到過，在訓練神經網絡時，最好是對一個batch的數據進行操作，同時還需要對數據進行shuffle和並行加速等。對此，PyTorch提供了DataLoader幫助我們實現這些功能。

 

DataLoader的函數定義如下： DataLoader(dataset, batch_size=1, shuffle=False, sampler=None, num_workers=0, collate_fn=default_collate, pin_memory=False, drop_last=False)

• dataset：加載的數據集(Dataset對象)
• batch_size：batch size
• shuffle:：是否將數據打亂
• sampler： 樣本抽樣，后續會詳細介紹
• num_workers：使用多進程加載的進程數，0代表不使用多進程
• collate_fn： 如何將多個樣本數據拼接成一個batch，一般使用默認的拼接方式即可
• pin_memory：是否將數據保存在pin memory區，pin memory中的數據轉到GPU會快一些
• drop_last：dataset中的數據個數可能不是batch_size的整數倍，drop_last為True會將多出來不足一個batch的數據丟棄

在數據處理中，有時會出現某個樣本無法讀取等問題，比如某張圖片損壞。這時在__getitem__函數中將出現異常，此時最好的解決方案即是將出錯的樣本剔除。如果實在是遇到這種情況無法處理，則可以返回None對象，然后在Dataloader中實現自定義的collate_fn，將空對象過濾掉。但要注意，在這種情況下dataloader返回的batch數目會少於batch_size。

class NewDogCat(DogCat): # 繼承前面實現的DogCat數據集
    def __getitem__(self, index):
        try:
            # 調用父類的獲取函數，即 DogCat.__getitem__(self, index)
            return super(NewDogCat,self).__getitem__(index)
        except:
            return None, None

from torch.utils.data.dataloader import default_collate # 導入默認的拼接方式
def my_collate_fn(batch):
    '''
    batch中每個元素形如(data, label)
    '''
    # 過濾為None的數據
    batch = list(filter(lambda x:x[0] is not None, batch))
    if len(batch) == 0: return t.Tensor()
    return default_collate(batch) # 用默認方式拼接過濾后的batch數據

dataloader = DataLoader(dataset, 2, collate_fn=my_collate_fn, num_workers=1,shuffle=True)
for batch_datas, batch_labels in dataloader:
    print(batch_datas.size(),batch_labels.size())

對於諸如樣本損壞或數據集加載異常等情況，還可以通過其它方式解決。例如但凡遇到異常情況，就隨機取一張圖片代替：

class NewDogCat(DogCat): def __getitem__(self, index): try: return super(NewDogCat, self).__getitem__(index) except: new_index = random.randint(0, len(self)-1) return self[new_index]

DataLoader里面並沒有太多的魔法方法，它封裝了Python的標准庫multiprocessing，使其能夠實現多進程加速。在此提幾點關於Dataset和DataLoader使用方面的建議：

1. 高負載的操作放在__getitem__中，如加載圖片等。
2. dataset中應盡量只包含只讀對象，避免修改任何可變對象，利用多線程進行操作。

第一點是因為多進程會並行的調用__getitem__函數，將負載高的放在__getitem__函數中能夠實現並行加速。 第二點是因為dataloader使用多進程加載，如果在Dataset實現中使用了可變對象，可能會有意想不到的沖突。在多線程/多進程中，修改一個可變對象，需要加鎖，但是dataloader的設計使得其很難加鎖（在實際使用中也應盡量避免鎖的存在），因此最好避免在dataset中修改可變對象。例如下面就是一個不好的例子，在多進程處理中self.num可能與預期不符，這種問題不會報錯，因此難以發現。如果一定要修改可變對象，建議使用Python標准庫Queue中的相關數據結構。

class BadDataset(Dataset):
    def __init__(self):
        self.datas = range(100)
        self.num = 0 # 取數據的次數
    def __getitem__(self, index):
        self.num += 1
        return self.datas[index]

使用Python multiprocessing庫的另一個問題是，在使用多進程時，如果主程序異常終止（比如用Ctrl+C強行退出），相應的數據加載進程可能無法正常退出。這時你可能會發現程序已經退出了，但GPU顯存和內存依舊被占用着，或通過top、ps aux依舊能夠看到已經退出的程序，這時就需要手動強行殺掉進程。建議使用如下命令：

ps x | grep <cmdline> | awk '{print $1}' | xargs kill

PyTorch中還單獨提供了一個sampler模塊，用來對數據進行采樣。常用的有隨機采樣器：RandomSampler，當dataloader的shuffle參數為True時，系統會自動調用這個采樣器，實現打亂數據。默認的是采用SequentialSampler，它會按順序一個一個進行采樣。這里介紹另外一個很有用的采樣方法： WeightedRandomSampler，它會根據每個樣本的權重選取數據，在樣本比例不均衡的問題中，可用它來進行重采樣。

構建WeightedRandomSampler時需提供兩個參數：每個樣本的權重weights、共選取的樣本總數num_samples，以及一個可選參數replacement。權重越大的樣本被選中的概率越大，待選取的樣本數目一般小於全部的樣本數目。replacement用於指定是否可以重復選取某一個樣本，默認為True，即允許在一個epoch中重復采樣某一個數據。如果設為False，則當某一類的樣本被全部選取完，但其樣本數目仍未達到num_samples時，sampler將不會再從該類中選擇數據，此時可能導致weights參數失效。下面舉例說明。

dataset = DogCat('data/dogcat/', transforms=transform)

# 狗的圖片被取出的概率是貓的概率的兩倍
# 兩類圖片被取出的概率與weights的絕對大小無關，只和比值有關
weights = [2 if label == 1 else 1 for data, label in dataset]
weights

from torch.utils.data.sampler import  WeightedRandomSampler
sampler = WeightedRandomSampler(weights,\
                                num_samples=9,\
                                replacement=True)
dataloader = DataLoader(dataset,
                        batch_size=3,
                        sampler=sampler)
for datas, labels in dataloader:
    print(labels.tolist())

 

計算機視覺工具包：torchvision

torchvision主要包含三部分：

• models：提供深度學習中各種經典網絡的網絡結構以及預訓練好的模型，包括AlexNet、VGG系列、ResNet系列、Inception系列等。
• datasets： 提供常用的數據集加載，設計上都是繼承torhc.utils.data.Dataset，主要包括MNIST、CIFAR10/100、ImageNet、COCO等。
• transforms：提供常用的數據預處理操作，主要包括對Tensor以及PIL Image對象的操作。

from torchvision import models
from torch import nn
# 加載預訓練好的模型，如果不存在會進行下載
# 預訓練好的模型保存在 ~/.torch/models/下面
resnet34 = models.squeezenet1_1(pretrained=True, num_classes=1000)

# 修改最后的全連接層為10分類問題（默認是ImageNet上的1000分類）
resnet34.fc=nn.Linear(512, 10)


from torchvision import datasets
# 指定數據集路徑為data，如果數據集不存在則進行下載
# 通過train=False獲取測試集
dataset = datasets.MNIST('data/', download=True, train=False, transform=transform)


from torchvision import transforms 
to_pil = transforms.ToPILImage()
to_pil(t.randn(3, 64, 64))

torchvision還提供了兩個常用的函數。一個是make_grid，它能將多張圖片拼接成一個網格中；另一個是save_img，它能將Tensor保存成圖片。

dataloader = DataLoader(dataset, shuffle=True, batch_size=16)
from torchvision.utils import make_grid, save_image
dataiter = iter(dataloader)
img = make_grid(next(dataiter)[0], 4) # 拼成4*4網格圖片，且會轉成３通道
to_img(img)

save_image(img, 'a.png')

 

可視化工具

在訓練神經網絡時，我們希望能更直觀地了解訓練情況，包括損失曲線、輸入圖片、輸出圖片、卷積核的參數分布等信息。這些信息能幫助我們更好地監督網絡的訓練過程，並為參數優化提供方向和依據。最簡單的辦法就是打印輸出，但其只能打印數值信息，不夠直觀，同時無法查看分布、圖片、聲音等。在本節，我們將介紹兩個深度學習中常用的可視化工具：Tensorboard和Visdom。

Tensorboard

Tensorboard也是一個相對獨立的工具，只要用戶保存的數據遵循相應的格式，tensorboard就能讀取這些數據並進行可視化。這里我們將主要介紹如何在PyTorch中使用tensorboardX1進行訓練損失的可視化。 TensorboardX是將Tensorboard的功能抽取出來，使得非TensorFlow用戶也能使用它進行可視化，幾乎支持原生TensorBoard的全部功能。

tensorboardX的使用非常簡單。首先用如下命令啟動tensorboard：

下面舉例說明tensorboardX的使用。

tensorboard --logdir <your/running/dir> --port <your_bind_port> # 啟動tensorboard


from tensorboardX import SummaryWriter
# 構建logger對象，logdir用來指定log文件的保存路徑
# flush_secs用來指定刷新同步間隔
logger = SummaryWriter(log_dir='experimient_cnn', flush_secs=2)
for ii in range(100):
    logger.add_scalar('data/loss', 10-ii**0.5)
    logger.add_scalar('data/accuracy', ii**0.5/10)

打開瀏覽器輸入http://localhost:6006（其中6006應改成你的tensorboard所綁定的端口），即可看到如圖2所示的結果。圖2: tensorboard可視化結果左側的Horizontal Axis下有三個選項，分別是：

Step：根據步長來記錄，log_value時如果有步長，則將其作為x軸坐標描點畫線。
Relative：用前后相對順序描點畫線，可認為logger自己維護了一個step屬性，每調用一次log_value就自動加１。
Wall：按時間排序描點畫線。
左側的Smoothing條可以左右拖動，用來調節平滑的幅度。點擊右上角的刷新按鈕可立即刷新結果，默認是每30s自動刷新數據。可見tensorboard_logger的使用十分簡單，但它只能統計簡單的數值信息，不支持其它功能。

感興趣的讀者可以從github項目主頁獲取更多信息，本節將把更多的內容留給另一個可視化工具：Visdom。

 

Visdom

Visdom可以創造、組織和共享多種數據的可視化，包括數值、圖像、文本，甚至是視頻，其支持PyTorch、Torch及Numpy。用戶可通過編程組織可視化空間，或通過用戶接口為生動數據打造儀表板，檢查實驗結果或調試代碼。

Visdom中有兩個重要概念：

• env：環境。不同環境的可視化結果相互隔離，互不影響，在使用時如果不指定env，默認使用main。不同用戶、不同程序一般使用不同的env。
• pane：窗格。窗格可用於可視化圖像、數值或打印文本等，其可以拖動、縮放、保存和關閉。一個程序中可使用同一個env中的不同pane，每個pane可視化或記錄某一信息。

 

cuda:

在PyTorch中以下數據結構分為CPU和GPU兩個版本：

• Tensor
• nn.Module（包括常用的layer、loss function，以及容器Sequential等）

它們都帶有一個.cuda方法，調用此方法即可將其轉為對應的GPU對象。注意，tensor.cuda會返回一個新對象，這個新對象的數據已轉移至GPU，而之前的tensor還在原來的設備上（CPU）。而module.cuda則會將所有的數據都遷移至GPU，並返回自己。所以module = module.cuda()和module.cuda()所起的作用一致。

nn.Module在GPU與CPU之間的轉換，本質上還是利用了Tensor在GPU和CPU之間的轉換。nn.Module的cuda方法是將nn.Module下的所有parameter（包括子module的parameter）都轉移至GPU，而Parameter本質上也是tensor(Tensor的子類)。

下面將舉例說明，這部分代碼需要你具有兩塊GPU設備。

P.S. 為什么將數據轉移至GPU的方法叫做.cuda而不是.gpu，就像將數據轉移至CPU調用的方法是.cpu？這是因為GPU的編程接口采用CUDA，而目前並不是所有的GPU都支持CUDA，只有部分Nvidia的GPU才支持。PyTorch未來可能會支持AMD的GPU，而AMD GPU的編程接口采用OpenCL，因此PyTorch還預留着.cl方法，用於以后支持AMD等的GPU。

 

# 不指定所使用的GPU設備，將默認使用第1塊GPU
tensor = tensor.cuda()
tensor.is_cuda # True

module = nn.Linear(3, 4)
module.cuda(device = 1)
module.weight.is_cuda # True

class VeryBigModule(nn.Module):
    def __init__(self):
        super(VeryBigModule, self).__init__()
        self.GiantParameter1 = t.nn.Parameter(t.randn(100000, 20000)).cuda(0)
        self.GiantParameter2 = t.nn.Parameter(t.randn(20000, 100000)).cuda(1)
    
    def forward(self, x):
        x = self.GiantParameter1.mm(x.cuda(0))
        x = self.GiantParameter2.mm(x.cuda(1))
        return x

上面最后一部分中，兩個Parameter所占用的內存空間都非常大，大概是8個G，如果將這兩個都同時放在一塊GPU上幾乎會將顯存占滿，無法再進行任何其它運算。此時可通過這種方式將不同的計算分布到不同的GPU中。

 

關於使用GPU的一些建議：

• GPU運算很快，但對於很小的運算量來說，並不能體現出它的優勢，因此對於一些簡單的操作可直接利用CPU完成
• 數據在CPU和GPU之間，以及GPU與GPU之間的傳遞會比較耗時，應當盡量避免
• 在進行低精度的計算時，可以考慮HalfTensor，它相比於FloatTensor能節省一半的顯存，但需千萬注意數值溢出的情況。

 另外這里需要專門提一下，大部分的損失函數也都屬於nn.Moudle，但在使用GPU時，很多時候我們都忘記使用它的.cuda方法，這在大多數情況下不會報錯，因為損失函數本身沒有可學習的參數（learnable parameters）。但在某些情況下會出現問題，為了保險起見同時也為了代碼更規范，應記得調用criterion.cuda。下面舉例說明

# 交叉熵損失函數，帶權重
criterion = t.nn.CrossEntropyLoss(weight=t.Tensor([1, 3]))
input = t.randn(4, 2).cuda()
target = t.Tensor([1, 0, 0, 1]).long().cuda()

# 下面這行會報錯，因weight未被轉移至GPU
# loss = criterion(input, target)

# 這行則不會報錯
criterion.cuda()
loss = criterion(input, target)

criterion._buffers

而除了調用對象的.cuda方法之外，還可以使用torch.cuda.device，來指定默認使用哪一塊GPU，或使用torch.set_default_tensor_type使程序默認使用GPU，不需要手動調用cuda。

# 如果未指定使用哪塊GPU，默認使用GPU 0
x = t.cuda.FloatTensor(2, 3)
# x.get_device() == 0
y = t.FloatTensor(2, 3).cuda()
# y.get_device() == 0

# 指定默認使用GPU 1
with t.cuda.device(1):    
    # 在GPU 1上構建tensor
    a = t.cuda.FloatTensor(2, 3)

    # 將tensor轉移至GPU 1
    b = t.FloatTensor(2, 3).cuda()
    print(a.get_device() == b.get_device() == 1 )

    c = a + b
    print(c.get_device() == 1)

    z = x + y
    print(z.get_device() == 0)

    # 手動指定使用GPU 0
    d = t.randn(2, 3).cuda(0)
    print(d.get_device() == 2)

如果服務器具有多個GPU，tensor.cuda()方法會將tensor保存到第一塊GPU上，等價於tensor.cuda(0)。此時如果想使用第二塊GPU，需手動指定tensor.cuda(1)，而這需要修改大量代碼，很是繁瑣。這里有兩種替代方法：

• 一種是先調用t.cuda.set_device(1)指定使用第二塊GPU，后續的.cuda()都無需更改，切換GPU只需修改這一行代碼。
• 更推薦的方法是設置環境變量CUDA_VISIBLE_DEVICES，例如當export CUDA_VISIBLE_DEVICE=1（下標是從0開始，1代表第二塊GPU），只使用第二塊物理GPU，但在程序中這塊GPU會被看成是第一塊邏輯GPU，因此此時調用tensor.cuda()會將Tensor轉移至第二塊物理GPU。CUDA_VISIBLE_DEVICES還可以指定多個GPU，如export CUDA_VISIBLE_DEVICES=0,2,3，那么第一、三、四塊物理GPU會被映射成第一、二、三塊邏輯GPU，tensor.cuda(1)會將Tensor轉移到第三塊物理GPU上。

設置CUDA_VISIBLE_DEVICES有兩種方法，一種是在命令行中CUDA_VISIBLE_DEVICES=0,1 python main.py，一種是在程序中import os;os.environ["CUDA_VISIBLE_DEVICES"] = "2"。如果使用IPython或者Jupyter notebook，還可以使用%env CUDA_VISIBLE_DEVICES=1,2來設置環境變量。

從 0.4 版本開始，pytorch新增了tensor.to(device)方法，能夠實現設備透明，便於實現CPU/GPU兼容。這部份內容已經在第三章講解過了。

 

從PyTorch 0.2版本中，PyTorch新增分布式GPU支持。分布式是指有多個GPU在多台服務器上，而並行一般指的是一台服務器上的多個GPU。分布式涉及到了服務器之間的通信，因此比較復雜，PyTorch封裝了相應的接口，可以用幾句簡單的代碼實現分布式訓練。分布式對普通用戶來說比較遙遠，因為搭建一個分布式集群的代價十分大，使用也比較復雜。相比之下一機多卡更加現實。對於分布式訓練，這里不做太多的介紹，感興趣的讀者可參考文檔1

 

單機多卡並行

要實現模型單機多卡十分容易，直接使用 new_module = nn.DataParallel(module, device_ids), 默認會把模型分布到所有的卡上。多卡並行的機制如下：

• 將模型（module）復制到每一張卡上
• 將形狀為（N,C,H,W）的輸入均等分為 n份（假設有n張卡），每一份形狀是（N/n, C,H,W）,然后在每張卡前向傳播，反向傳播，梯度求平均。要求batch-size 大於等於卡的個數(N>=n)

 

持久化

在PyTorch中，以下對象可以持久化到硬盤，並能通過相應的方法加載到內存中：

• Tensor
• Variable
• nn.Module
• Optimizer

本質上上述這些信息最終都是保存成Tensor。Tensor的保存和加載十分的簡單，使用t.save和t.load即可完成相應的功能。在save/load時可指定使用的pickle模塊，在load時還可將GPU tensor映射到CPU或其它GPU上。

我們可以通過t.save(obj, file_name)等方法保存任意可序列化的對象，然后通過obj = t.load(file_name)方法加載保存的數據。對於Module和Optimizer對象，這里建議保存對應的state_dict，而不是直接保存整個Module/Optimizer對象。Optimizer對象保存的主要是參數，以及動量信息，通過加載之前的動量信息，能夠有效地減少模型震盪，下面舉例說明。

a = t.Tensor(3, 4)
if t.cuda.is_available():
        a = a.cuda(1) # 把a轉為GPU1上的tensor,
        t.save(a,'a.pth')
        
        # 加載為b, 存儲於GPU1上(因為保存時tensor就在GPU1上)
        b = t.load('a.pth')
        
        # 加載為c, 存儲於CPU
        c = t.load('a.pth', map_location=lambda storage, loc: storage)
        
        # 加載為d, 存儲於GPU0上
        d = t.load('a.pth', map_location={'cuda:1':'cuda:0'})

t.set_default_tensor_type('torch.FloatTensor')
from torchvision.models import SqueezeNet
model = SqueezeNet()
# module的state_dict是一個字典
model.state_dict().keys()


# Module對象的保存與加載
t.save(model.state_dict(), 'squeezenet.pth')
model.load_state_dict(t.load('squeezenet.pth'))

all_data = dict(
    optimizer = optimizer.state_dict(),
    model = model.state_dict(),
    info = u'模型和優化器的所有參數'
)
t.save(all_data, 'all.pth')

all_data = t.load('all.pth')
all_data.keys()