def main(args):
#########################################
### Loading Dataset ###
#########################################
if args.prior_dist == "G":
transform = transforms.Compose([transforms.Resize([64, 64], 1),
transforms.ToTensor(),
transforms.Normalize([.5], [.5])])
elif args.prior_dist == "B":
transform = transforms.Compose([trainsforms.Resize([64, 64], 1),
transforms.ToTensor()])
target_transform = copy.deepcopy(transform)
transform.transforms.insert(0, transforms.RandomAffine(degrees=10, translate=(0.1, 0.1)))
transform.transforms.insert(0, transforms.ColorJitter(0.4, 0.4, 0.4))
print(transform.transforms)
train_root = os.path.join(args.data_path, "train")
train_data, train_loader = load_data(train_root, args.batch_size,
transform, target_transform,
input_mode=args.input_mode,
condition=True if args.net == "CVAE" else False)
test_root = os.path.join(args.data_path, "test")
_, test_loader = load_data(test_root, args.batch_size,
transform, target_transform,
input_mode=args.input_mode,
condition=True if args.net == "CVAE" else False)
if args.validate:
vali_root = os.path.join(args.data_path, "validation")
_, vali_loader = load_data(vali_root, args.batch_size,
transform, target_transform,
input_mode=args.input_mode,
condition=True if args.net == "CVAE" else False)
else:
vali_loader = None
img_size=train_data[0][0].shape
#########################################
### Building Model ###
#########################################
model = Model(args)
model.build(eval(args.net), img_size=img_size, z_dim=args.z_dim)
if args.model_load_path:
model_name = "*.pt"
pt_file = load_model(args.model_load_path, model_name)
model.load(pt_file)
if args.mode == "train":
model.train(train_loader, vali_loader)
elif args.mode == "test":
result_img = model.test(test_loader)
result = data + [result_img]
if args.prior_dist == "G":
grid_img = make_result_img(result, normalize=True, range=(-1., 1.))
else:
grid_img = make_result_img(result)
utils.save_image(grid_img, "{}{}_test.png".format(args.log_path, args.net))
def train(self, train_loader, vali_loader=None):
self.net.train()
for epoch in range(1, self.epoch + 1):
total_sum, rec_sum, kld_sum = .0, .0, .0
for i, batch in enumerate(train_loader, 1):
batch_inputs = [b.to(self.device) for b in batch[:-1]]
batch_target = batch[-1].to(self.device)
self.optim.zero_grad()
logits, outputs, mu, log_var = self.net(*batch_inputs)
if self.prior_dist == "G":
total_loss, rec_loss, kld_loss = self.loss(
outputs, batch_target, mu, log_var)
elif self.prior_dist == "B":
total_loss, rec_loss, kld_loss = self.loss(
logits, batch_target, mu, log_var)
total_sum += total_loss.item()
rec_sum += rec_loss.item()
kld_sum += kld_loss.item()
total_loss.backward()
self.optim.step()
self.scheduler.step()
self.loss_list['total_loss'] = total_sum / i
self.loss_list['rec_loss'] = rec_sum / i
self.loss_list['kld_loss'] = kld_sum / i
if epoch % self.log_iter == 0:
if vali_loader is not None:
self.test(vali_loader)
data = iter(train_loader).next()
imgs = [b.to(self.device) for b in data[:-1]]
with torch.no_grad():
logits, outputs, _, _ = self.net(*imgs)
imgs = [b[:self.rec_num] for b in data]
if self.prior_dist == "G":
imgs.append(outputs.detach().cpu()[:self.rec_num])
self.result_img = make_result_img(imgs,
normalize=True,
range=(-1., 1.))
elif self.prior_dist == "B":
imgs.append(
torch.sigmoid(logits.detach().cpu())[:self.rec_num])
self.result_img = make_result_img(imgs)
self.logger(epoch)
if epoch % self.model_save_iter == 0:
torch.save(
self.net.state_dict(), self.model_save_path +
"{}_{:0>6}.pt".format(self.name, epoch))
def main(args):
transform = transforms.Compose([
transforms.Resize([64, 64], 1),
transforms.ToTensor(),
transforms.Normalize([.5], [.5])
])
dataset = datasets.ImageFolder(root=args.data_path, transform=transform)
data = []
for i, j in dataset:
data.append(i)
data = torch.stack(data)
cvae = CVAE(img_size=data.shape[1:], z_dim=args.z_dim)
cvae.eval()
pt_files = glob.glob(os.path.join(args.model_load_path, "*.pt"))
pt_files.sort()
#data_1 = data[0][None, :, :, :].repeat([20, 1, 1, 1])
#data_2 = data[1][None, :, :, :].repeat([20, 1, 1, 1])
for i in range(len(pt_files)):
print(pt_files[i])
cvae.load_state_dict(torch.load(pt_files[i]))
#z_data = [torch.randn(data.shape[0], args.z_dim), data]
z_data = [
torch.randn(32, args.z_dim), data[None, 0].repeat([32, 1, 1, 1])
]
_, rec_img = cvae.decoder(*z_data)
grid_img = make_result_img(
[data[None, 0].repeat([32, 1, 1, 1]), rec_img],
normalize=True,
range=(-1., 1.))
utils.save_image(
grid_img, "{}CVAE_gen_result_{:0>2d}.png".format(args.log_path, i))
def main(args):
transform = transforms.Compose([
transforms.Resize([64, 64], 1),
transforms.ToTensor(),
transforms.Normalize([.5], [.5])
])
vae = VAE(img_size=[3, 64, 64], z_dim=args.z_dim)
pt_file = load_model(args.model_load_path, "*.pt")
vae.load_state_dict(torch.load(pt_file))
vae.eval()
if args.decode:
#z = torch.randn(16, args.z_dim)
z = torch.from_numpy(
np.loadtxt("./denorm_output_log.txt",
delimiter=",",
dtype=np.float32)[:, 5:])
#for i in range(10):
for i in range(z.shape[0]):
#z_ = z.repeat(50, 1)
#z_[:, i] = torch.linspace(-4, 4, 50)
logits, result = vae.decoder(z[i, None])
grid_img = make_result_img([result],
normalize=True,
range=(-1., 1.))
utils.save_image(grid_img,
"./rec/VAE_decode_result_{:0>6d}.png".format(i))
if args.rec_img:
dataset = datasets.ImageFolder(root=args.data_path,
transform=transform)
data, label = [], []
for i, j in dataset:
data.append(i)
label.append(j)
data = torch.stack(data)
label = torch.tensor(label)
_, result, _, _ = vae(data)
grid_img = make_result_img([data, result],
normalize=True,
range=(-1., 1.))
utils.save_image(grid_img,
"{}VAE_reconstract_result.png".format(args.log_path))
def main(args):
transform = transforms.Compose([
transforms.Resize([64, 64], 1),
transforms.ToTensor(),
transforms.Normalize([.5], [.5])
])
dataset = datasets.ImageFolder(root=args.data_path, transform=transform)
data = []
for i, j in dataset:
data.append(i)
data = torch.stack(data)
cvae = CVAE(img_size=data.shape[1:], z_dim=args.z_dim)
pt_file = load_model(args.model_load_path, "*10.pt")
cvae.load_state_dict(torch.load(pt_file))
cvae.eval()
if args.decode:
#z_data = [torch.randn(data.shape[0], args.z_dim), data]
#_, rec_img = cvae.decoder(*z_data)
#grid_img = make_result_img([data, rec_img], normalize=True, range=(-1., 1.))
#utils.save_image(grid_img, "{}CVAE_gen_result.png".format(args.log_path))
cond_img = data[None, 0].repeat([32, 1, 1, 1])
z_data = [torch.randn(cond_img.shape[0], args.z_dim), cond_img]
_, rec_img = cvae.decoder(*z_data)
#grid_img = make_result_img([rec_img], normalize=True, range=(-1., 1.))
#utils.save_image(grid_img, "{}CVAE_gen_result_same_cond.png".format(args.log_path))
for i in range(rec_img.shape[0]):
utils.save_image(rec_img[i],
"{}goal{:0>6d}.png".format(args.log_path, i),
normalize=True,
range=(-1., 1.))
if args.gen_seq:
for i, d in enumerate(data):
cond_img = data[None, i]
z_data = [torch.randn(1, args.z_dim), cond_img]
_, rec_img = cvae.decoder(*z_data)
grid_img = make_result_img([cond_img, rec_img],
normalize=True,
range=(-1., 1.))
utils.save_image(
grid_img, "{}res_state-goal/CVAE_gen_{:0>6d}.png".format(
args.log_path, i))
utils.save_image(rec_img,
"{}res_goal/CVAE_gen_{:0>6d}.png".format(
args.log_path, i),
normalize=True,
range=(-1., 1.))
def main(args):
transform = transforms.Compose([
transforms.Resize([64, 64], 1),
transforms.ToTensor(),
transforms.Normalize([.5], [.5])
])
vae = VAE(img_size=[3, 64, 64], z_dim=args.z_dim)
pt_file = load_model(args.model_load_path, "*.pt")
vae.load_state_dict(torch.load(pt_file))
vae.eval()
z = np.loadtxt("../denorm_pb_list.txt", delimiter=",", dtype=np.float32)
x = np.linspace(z[0], z[5], 100)
print(x.shape)
_, result = vae.decoder(torch.from_numpy(x))
grid_img = make_result_img([result], normalize=True, range=(-1., 1.))
utils.save_image(grid_img, "./x.png")
for sample in data.samples:
goal_data[-1].append((sample[0], i))
print("=" * 5, "Loading cond data", "=" * 5)
dataset = CondDataSet(env_data,
goal_data,
transform=transform,
target_transform=target_transform,
input_mode=input_mode)
else:
print("=" * 5, "Loading data", "=" * 5)
dataset = DataSet(env_data,
transform=transform,
target_transform=target_transform)
img_size = dataset[0][0].shape
data_loader = DataLoader(dataset,
batch_size=batch_size,
shuffle=True,
num_workers=8)
return dataset, data_loader
if __name__ == "__main__":
dataset, data_loader = load_data("./data/20200604/train/", 16,
transforms.ToTensor(),
transforms.ToTensor())
x = iter(data_loader).next()
grid_img = make_result_img(x[:-1])
utils.save_image(grid_img, "./grid_img.png")