collate_fn 參數(shù)的目的主要是為了隨心所欲的轉(zhuǎn)變數(shù)據(jù)的類型,這個(gè)數(shù)據(jù)是用DataLoader加載的,比如img，target,下面這篇文章主要給大家介紹了關(guān)于Pytorch使用技巧之Dataloader中的collate_fn參數(shù)的相關(guān)資料,需要的朋友可以參考下

以MNIST為例

from torchvision import datasets
mnist = datasets.MNIST(root='./data/', train=True, download=True)
print(mnist[0])

結(jié)果

(<PIL.Image.Image image mode=L size=28x28 at 0x196E3F1D898>, 5)

MINIST數(shù)據(jù)集的dataset是由一張圖片和一個(gè)label組成的元組

dataloader = torch.utils.data.DataLoader(dataset=mnist, batch_size=2, shuffle=True,collate_fn=lambda x:x)
for each in dataloader:
    print(each)
    break

結(jié)果

[(<PIL.Image.Image image mode=L size=28x28 at 0x2CB3B105630>, 0), (<PIL.Image.Image image mode=L size=28x28 at 0x2CB3B105668>, 2)]

collate_fn為lamda x:x時(shí)表示對(duì)傳入進(jìn)來的數(shù)據(jù)不做處理

下面自定義collate_fn看看什么效果

def collate(data):
    img = []
    label = []
    for each in data:
        img.append(each[0])
        label.append(each[1])
    return img,label
dataloader = torch.utils.data.DataLoader(dataset=mnist, batch_size=2, shuffle=True,collate_fn=lambda x:collate(x))
for each in dataloader:
    print(each)
    break

結(jié)果

([<PIL.Image.Image image mode=L size=28x28 at 0x241433A36D8>, <PIL.Image.Image image mode=L size=28x28 at 0x241433A3710>], [9, 3])

說明：若不設(shè)置collate_fn參數(shù)則會(huì)使用默認(rèn)處理函數(shù)

但必須保證傳進(jìn)來的數(shù)據(jù)都是tensor格式否則會(huì)報(bào)錯(cuò)

附：DataLoader完整的參數(shù)表如下：

class torch.utils.data.DataLoader(
    dataset,
    batch_size=1,
    shuffle=False,
    sampler=None,
    batch_sampler=None,
    num_workers=0,
    collate_fn=<function default_collate>,
    pin_memory=False,
    drop_last=False,
    timeout=0,
    worker_init_fn=None)

DataLoader在數(shù)據(jù)集上提供單進(jìn)程或多進(jìn)程的迭代器

幾個(gè)關(guān)鍵的參數(shù)意思：

- shuffle：設(shè)置為True的時(shí)候，每個(gè)世代都會(huì)打亂數(shù)據(jù)集

- collate_fn：如何取樣本的，我們可以定義自己的函數(shù)來準(zhǔn)確地實(shí)現(xiàn)想要的功能

- drop_last：告訴如何處理數(shù)據(jù)集長(zhǎng)度除于batch_size余下的數(shù)據(jù)。True就拋棄，否則保留

總結(jié)

到此這篇關(guān)于Pytorch使用技巧之Dataloader中的collate_fn參數(shù)的文章就介紹到這了,更多相關(guān)Dataloader中的collate_fn參數(shù)內(nèi)容請(qǐng)搜索腳本之家以前的文章或繼續(xù)瀏覽下面的相關(guān)文章希望大家以后多多支持腳本之家！

您可能感興趣的文章: