data_loader.py 文件源码

def __init__(self, path, img_size, batch_size, is_cuda):
        self._img_files = os.listdir(path)
        self._path = path
        self._is_cuda = is_cuda
        self._step = 0
        self._batch_size = batch_size
        self.sents_size = len(self._img_files)
        self._stop_step = self.sents_size // batch_size

        self._encode = transforms.Compose([
                            transforms.Scale(img_size),
                            transforms.RandomCrop(img_size),
                            transforms.ToTensor()
                        ])