def load_sprite(params):
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Lambda(to_float),
])
trainset = datasets.Sprites(params["data_dir"],
rnd_background=params["rnd_bkg"],
train=True,
transform=transform,
n=params["num_samples"],
max_num_objs=params["max_num_objs"],
min_num_objs=params["min_num_objs"])
testset = datasets.Sprites(params["data_dir"],
rnd_background=params["rnd_bkg"],
train=False,
transform=transform,
n=params["num_samples"],
max_num_objs=params["max_num_objs"],
min_num_objs=params["min_num_objs"])
trainloader = torch.utils.data.DataLoader(trainset,
batch_size=params["batch_size"],
shuffle=True,
num_workers=0)
testloader = torch.utils.data.DataLoader(testset,
batch_size=params["batch_size"],
shuffle=True,
num_workers=0)
return trainloader, testloader
def sprite_experiment():
conf = config.sprite_config
transform = transforms.Compose([transforms.ToTensor(),
transforms.Lambda(lambda x: x.float()),
])
trainset = datasets.Sprites(conf.data_dir, train=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset,
batch_size=conf.batch_size // conf.subdivs,
shuffle=True, num_workers=2)
monet = model.Monet(conf, 64, 64).cuda()
if conf.parallel:
monet = nn.DataParallel(monet)
run_training(monet, conf, trainloader)