I want to train a model on my own dataset, I implement a custom dataset and a custom dataloader according to Data Loading and Processing Tutorial. When I run the training process, I can not get the data patch. The process hang up for a long time and no errors are reported.
Dataset
class CustomDataset(BaseDataset):
def initialize(self, opt):
self.root = opt.dataroot
print 'self.root: {}'.format(self.root)
self.imgs = os.listdir(self.root)
self.imgs = sorted(self.imgs)
self.fdir = os.path.join(os.path.abspath(os.path.join(self.root,'..')), 'features')
self.feas = [os.path.join(self.fdir, img) for img in self.imgs]
self.imgs = [os.path.join(self.root, img) for img in self.imgs]
self.transform = transforms.Compose([transforms.ToTensor(),])
def __getitem__(self, index):
i_path = self.imgs[index]
f_path = self.feas[index]
img = Image.open(i_path).convert('RGB')
img = self.transform(img)
feature = torch.Tensor(cPickle.load(open(f_path, 'rb')))
#feature = torch.Tensor([1,])
#print 'len(img):{}'.format(img.size())
#print 'len(feature):{}'.format(feature.size())
input_dict = {'img': img, 'feature': feature}
return input_dict
def __len__(self):
return len(self.imgs)
def name(self):
return 'CustomDataset'
DataLoader
def CreateDataset(opt):
dataset = None
from custom_dataset import CustomDataset
dataset = CustomDataset()
dataset.initialize(opt)
return dataset
class CustomDataLoader(BaseDataLoader):
def name(self):
return 'CustomDataLoader'
def initialize(self, opt):
BaseDataLoader.initialize(self, opt)
self.dataset = CreateDataset(opt)
self.dataloader = torch.utils.data.DataLoader(
self.dataset,
batch_size=opt.batchSize,
shuffle=not opt.shuffle,
num_workers=int(opt.workers))
def load_data(self):
return self
def __len__(self):
return min(len(self.dataset), self.opt.max_dataset_size)
def __iter__(self):
for i, data in enumerate(self.dataloader):
print 'custom_dataloader:{}'.format(i)
if i >= self.opt.max_dataset_size:
break
yield data