def _make_batch_generator(self):
# data load and construct batch generator
self.logger.info("Creating dataset...")
trainset_list = []
for i in range(len(self.cfg.trainset)):
trainset_list.append(eval(self.cfg.trainset[i])("train"))
trainset_loader = DatasetLoader(trainset_list, True, transforms.Compose([\
transforms.ToTensor(),
transforms.Normalize(mean=cfg.pixel_mean, std=cfg.pixel_std)]\
))
batch_generator = DataLoader(dataset=trainset_loader,
batch_size=self.cfg.num_gpus *
self.cfg.batch_size,
shuffle=True,
num_workers=self.cfg.num_thread,
pin_memory=True)
self.joint_num = trainset_loader.joint_num[0]
self.itr_per_epoch = math.ceil(trainset_loader.__len__() /
cfg.num_gpus / cfg.batch_size)
self.batch_generator = batch_generator
def _make_batch_generator(self):
# data load and construct batch generator
self.logger.info("Creating dataset...")
testset = eval(self.cfg.testset)("test")
testset_loader = DatasetLoader(testset, False, transforms.Compose([\
transforms.ToTensor(),
transforms.Normalize(mean=cfg.pixel_mean, std=cfg.pixel_std)]\
))
batch_generator = DataLoader(dataset=testset_loader,
batch_size=self.cfg.num_gpus *
self.cfg.test_batch_size,
shuffle=False,
num_workers=self.cfg.num_thread,
pin_memory=True)
self.testset = testset
self.joint_num = testset_loader.joint_num
self.skeleton = testset_loader.skeleton
self.flip_pairs = testset.flip_pairs
self.tot_sample_num = testset_loader.__len__()
self.batch_generator = batch_generator
