pytorch学习笔记之DataLoader

叁歲伎倆 2022-12-20 11:11 277阅读 0赞

# Dataset & DataLoader基础 #

在PyTorch中一般需要使用DataLoader来小批量的加载数据，也就类似于mini-batch的操作，如下图所示：

![watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3FxXzI2ODIyMDI5_size_16_color_FFFFFF_t_70][]

（图片来自b站\[刘二大人\]上传视频《PyTorch深度学习实践》完结合集）

为此PyTorch提供了torch.utils.data.Dataset这样一个基类用于构建我们自己的数据集，我们要做的就是基于这个基类实现一个派生类，需要实现两个函数，如下代码框架所示：

import torch
    from torch.utils.data import Dataset
    from torch.utils.data import DataLoader
    
    class YourDataset(Dataset):
        def __init__(self):
            pass
    
        def __getitem__(self, index):
            """通过索引index获取对应的数据"""
            pass
    
        def __len__(self):
            """返回数据的条数"""
            pass
    
    dataset = YourDataset()
    train_loader = DataLoader(dataset=dataset, batch_size=32, shuffle=True, num_workers=2)
    
    for epoch in range(train_epoch):
        for i, data in enumerate(train_loader, 0):
            # do training

在实现上面的YourDataset类的时候有两种思路：

1.  如果数据量本身比较小的话，可以直接在\_\_init\_\_方法中将数据全部读到内存中，后面使用的时候直接在内存里面取；
2.  如果数据量较大的话用上面的策略就会出问题了，有很多种实现方法。例如，在\_\_init\_\_中仅仅保存数据的索引，在\_\_getitem\_\_中根据index和索引再去磁盘上加载数据；或者借助pandas的read\_csv方法中的iterators=True循环分批将数据读进内存。

如，我自己写的小批量数据集的Dataset代码：

import pandas as pd
    import torch
    import torch.nn as nn
    import torch.utils.data as data
    
    class DualInputModelDataLoader(data.Dataset):
        def __init__(self, file_path):
            self.min_sequence_len = 11
            self.max_sequence_len = 18
            self.AA_mapping = "ACDEFGHIKLMNPQRSTVWY-"
            
            df = pd.read_csv(file_path)
            self.data = self._parse_seq(df)
            self.len = len(self.data)
        
        def __getitem__(self, idx):
            return self.data[idx]
        
        def __len__(self):
            return self.len
        
        def _parse_seq(self, df):
            ret_list = []
            
            for _, row in df.iterrows():
                seq_a, seq_b, label = row['seq_a'], row['seq_b'], row['label']
                seq_a_len, seq_b_len = len(seq_a), len(seq_b)
                if seq_a_len < self.min_sequence_len or seq_a_len > self.max_sequence_len:
                    continue
                if seq_b_len < self.min_sequence_len or seq_b_len > self.max_sequence_len:
                    continue
                    
                padding_seq_a = self._sequence_padding(seq_a, 'middle', self.max_sequence_len)
                padding_seq_b = self._sequence_padding(seq_b, 'middle', self.max_sequence_len)
                ret_list.append((self._sequence_mapping(padding_seq_a), self._sequence_mapping(padding_seq_b), label))
            return ret_list
                
        def _sequence_padding(self, cdr3_sequence, padding_style, max_seq_len, padding_placeholder='-'):
            seq_len = len(cdr3_sequence)
    
            if padding_style == 'right':
                padding_seq = cdr3_sequence + padding_placeholder * (max_seq_len - seq_len)
            elif padding_style == 'left':
                padding_seq = padding_placeholder * (max_seq_len - seq_len) + cdr3_sequence
            elif padding_style == 'middle':
                padding_seq = cdr3_sequence[:seq_len // 2] + padding_placeholder * (max_seq_len - seq_len) + cdr3_sequence[seq_len // 2:]
            else:
                raise "error"
            return padding_seq
        
        def _sequence_mapping(self, seq):
            ret_list = [self.AA_mapping.index(aa) for aa in seq]
            return torch.LongTensor(ret_list)

[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3FxXzI2ODIyMDI5_size_16_color_FFFFFF_t_70]: /images/20221120/dace5bbd471a4d5f8a4204d15bf0b170.png