pytorch---之dataloader使用pandas读取csv(不一次性将csv全部读入内存)

最新推荐文章于 2025-05-19 16:50:36 发布

原创

最新推荐文章于 2025-05-19 16:50:36 发布 · 1.5w 阅读

24 ·

CC 4.0 BY-SA版权

我们使用pandas的read_csv()函数中的iterators=True 来将csv循环分批读取进内存，（如果不使用iterators参数的话，数据量很大的csv文件，会全部读到内存当中去，内存肯定不够）代码如下：

# -*- coding: utf-8 -*-

import csv
import pandas as pd
import numpy as np
import torch
import torch.utils.data as data


class FaceLandmarksDataset(data.Dataset):
	"""Face Landmarks dataset."""
	def __init__(self, csv_file):
		"""
        Args:
            csv_file (string): Path to the csv file with annotations.
            root_dir (string): Directory with all the images.
            transform (callable, optional): Optional transform to be applied
                on a sample.
        """
		self.landmarks_frame = pd.read_csv(csv_file, iterator=True)
	def __len__(self):
		#print len(self.landmarks_frame)
		#return len(self.landmarks_frame)
		return 1800000
	def __getitem__(self, idx):
		print idx
		landmarks = self.landmarks_frame.get_chunk(128).as_matrix().astype('float')
		# landmarks = self.landmarks_frame.ix[idx, 1:].