def main(args):
dim = 64
latent_size = 512
channels = 3
best_model_path = "./vae_lambda001.pth"
# should give from outside
output_path = args.output_path
# load model
model = VAE(d=dim, zsize=latent_size, channels=channels)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)
model.load_state_dict(torch.load(best_model_path, map_location=device))
# 4. sample from normal distribution
seed = 5
np.random.seed(seed)
samples = np.random.randn(32, latent_size)
torch_samples = torch.FloatTensor(samples)
torch_samples = torch_samples.to(device)
result = model.decode(torch_samples)
# concat images
rows = []
for row in range(8):
rows.append(
torch.cat([img for img in result[row * 4:(row + 1) * 4]], axis=1))
ret = torch.cat(rows, axis=2)
img = ret.cpu().detach().numpy().transpose(1, 2, 0)
scipy.misc.imsave(output_path, img)
def fake(config):
device = torch.device(
'cuda') if torch.cuda.is_available() else torch.device('cpu')
net = VAE().to(device)
net.load_state_dict((torch.load(config.model_path + config.model_name)))
net.eval()
with torch.no_grad():
fake_z = torch.randn((64, net.nz)).to(device)
fake_imgs = net.decode(fake_z).view(-1, 1, 28, 28).detach()
plt.figure(figsize=(8, 8))
plt.axis('off')
plt.title('Fake images')
plt.imshow(
np.transpose(
torchvision.utils.make_grid(fake_imgs,
padding=2,
normalize=True).cpu(), (1, 2, 0)))
plt.show()
def main(args):
args = parse_args()
# check if cuda available
# device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
# define dataset and dataloader
dataset = AirfoilDataset()
airfoil_x = dataset.get_x()
airfoil_dim = airfoil_x.shape[0]
airfoil_dataloader = DataLoader(dataset, batch_size=16, shuffle=True)
# hyperparameters
latent_dim = 16 # please do not change latent dimension
lr = args.lr # learning rate
num_epochs = args.num_epochs
# build the model
vae = VAE(airfoil_dim=airfoil_dim, latent_dim=latent_dim, lr=lr).to(device)
print("VAE model:\n", vae)
# define your loss function here
# loss = ?
# define optimizer for discriminator and generator separately
# train the VAE model
epoch_loss_arr = []
for epoch in range(num_epochs):
epoch_loss = 0.0
num_batches = 0
for n_batch, (local_batch, __) in enumerate(airfoil_dataloader):
num_batches += 1
y_real = local_batch.to(device)
# train VAE
y_real = torch.tensor(y_real).float().to(device)
y_gen, mean, logvar = vae(y_real)
var = torch.exp(logvar)
# calculate customized VAE loss
# loss = your_loss_func(...)
rc_loss = torch.dist(y_real, y_gen, 2)
prior_loss = -1.0 * torch.sum(
(logvar + 1.0), 1) + torch.sum(var, 1) + torch.sum(mean**2, 1)
loss = rc_loss + 1.0 * 0.5 * torch.mean(prior_loss)
vae.optimizer.zero_grad()
loss.backward()
vae.optimizer.step()
epoch_loss += loss.item()
# print loss while training
if (n_batch + 1) % 30 == 0:
print("Epoch: [{}/{}], Batch: {}, loss: {}".format(
epoch, num_epochs, n_batch, loss.item()))
epoch_loss_arr.append(epoch_loss / num_batches)
# test trained VAE model
num_samples = 100
# reconstuct airfoils
real_airfoils = dataset.get_y()[:num_samples]
recon_airfoils, __, __ = vae(torch.from_numpy(real_airfoils).to(device))
if 'cuda' in device:
recon_airfoils = recon_airfoils.detach().cpu().numpy()
else:
recon_airfoils = recon_airfoils.detach().numpy()
# randomly synthesize airfoils
noise = torch.randn(
(num_samples, latent_dim)).to(device) # create random noise
gen_airfoils = vae.decode(noise)
if 'cuda' in device:
gen_airfoils = gen_airfoils.detach().cpu().numpy()
else:
gen_airfoils = gen_airfoils.detach().numpy()
# plot real/reconstructed/synthesized airfoils
plot_airfoils(airfoil_x, real_airfoils, "vae_real_airfoils")
plot_airfoils(airfoil_x, recon_airfoils, "vae_reconstructed_airfoils")
plot_airfoils(airfoil_x, gen_airfoils, "vae_generated_airfoils")
plot_prop(epoch_loss_arr, "vae_train_loss_lr_{}".format(lr))
class FactorVAE(BaseVAE):
"""Class that implements the Factor Variational Auto-Encoder"""
def __init__(self, n_input, n_hidden, dim_z, n_output, gamma, binary=True, **kwargs):
"""initialize neural networks
:param gamma: weight for total correlation term in loss function
"""
super(FactorVAE, self).__init__()
self.dim_z = dim_z
self.binary = binary
self.gamma = gamma
self.input_size = (n_input,)
# VAE networks
self.vae = VAE(n_input, n_hidden, dim_z, n_output, binary, **kwargs)
# discriminator layers
D_hidden_num = 3
D_hidden_dim = 1000
D_hidden_dims = [D_hidden_dim] * D_hidden_num
D_act = nn.LeakyReLU
D_act_args = {"negative_slope": 0.2, "inplace": False}
D_output_dim = 2
self.discriminator = nns.create_mlp(self.dim_z, D_hidden_dims,
act_layer=D_act, act_args=D_act_args, norm=True)
self.discriminator = nn.Sequential(
self.discriminator,
nn.Linear(D_hidden_dim, D_output_dim))
def encode(self, x):
"""vae encode"""
return self.vae.encode(x)
def reparameterize(self, mu, logvar):
"""reparameterization trick"""
return self.vae.reparameterize(mu, logvar)
def decode(self, code):
"""vae deocde"""
return self.vae.decode(code)
def sample_latent(self, num, device, **kwargs):
"""vae sample latent"""
return self.vae.sample_latent(num, device, **kwargs)
def sample(self, num, device, **kwargs):
"""vae sample"""
return self.vae.sample(num, device, **kwargs)
def forward(self, input, no_dec=False):
"""autoencoder forward computation"""
encoded = self.encode(input)
mu, logvar = encoded
z = self.reparameterize(mu, logvar) # latent variable z
if no_dec:
# no decoding
return z
return self.decode(z), encoded, z
def decoded_to_output(self, decoded, **kwargs):
"""vae transform decoded result to output"""
return self.vae.decoded_to_output(decoded, **kwargs)
def reconstruct(self, input, **kwargs):
"""vae reconstruct"""
return self.vae.reconstruct(input, **kwargs)
def permute_dims(self, z):
"""permute separately each dimension of the z randomly in a batch
:param z: [B x D] tensor
:return: [B x D] tensor with each dim of D dims permuted randomly
"""
B, D = z.size()
# generate randomly permuted batch on each dimension
permuted = []
for i in range(D):
ind = torch.randperm(B)
permuted.append(z[:, i][ind].view(-1, 1))
return torch.cat(permuted, dim=1)
def loss_function(self, *inputs, **kwargs):
"""loss function described in the paper (eq. (2))"""
optim_part = kwargs['optim_part'] # the part to optimize
if optim_part == 'vae':
# update VAE
decoded = inputs[0]
encoded = inputs[1]
z = inputs[2]
x = inputs[3]
mu, logvar = encoded
# KL divergence term
KLD = -0.5 * (1 + logvar - mu.pow(2) - logvar.exp()).sum(1).mean()
if self.binary:
# likelihood term under Bernolli MLP decoder
MLD = F.binary_cross_entropy(decoded, x, reduction='sum').div(x.size(0))
else:
# likelihood term under Gaussian MLP decoder
mu_o, logvar_o = decoded
recon_x_distribution = Normal(
loc=mu_o, scale=torch.exp(0.5*logvar_o))
MLD = -recon_x_distribution.log_prob(x).sum(1).mean()
Dz = self.discriminator(z)
tc_loss = (Dz[:, :1] - Dz[:, 1:]).mean()
return {
"loss": KLD + MLD + self.gamma * tc_loss,
"KLD": KLD,
"MLD": MLD,
"tc_loss": tc_loss}
elif optim_part == 'discriminator':
# update discriminator
Dz = inputs[0]
Dz_pperm = inputs[1]
device = z.device
ones = torch.ones(
Dz.size(0), dtype=torch.long, requires_grad=False).to(device)
zeros = torch.zeros(
Dz.size(0), dtype=torch.long, requires_grad=False).to(device)
D_tc_loss = 0.5 * (F.cross_entropy(Dz, zeros) +
F.cross_entropy(Dz_pperm, ones))
return {"loss": D_tc_loss, "D_tc_loss": D_tc_loss}
else:
raise Exception("no such network to optimize: {}".format(optim_part))
plt.clf()
plt.scatter(latent_z[:, 0],
latent_z[:, 1],
c=latent_y,
s=6,
cmap='rainbow',
edgecolors='k',
linewidth=0.25)
plt.axes().set_aspect('equal')
plt.show()
image_rows = tuple()
for yr in np.linspace(-4, 4, 10):
image_row = tuple()
for xr in np.linspace(-4, 4, 10):
x_reconstructed = vae.decode(np.array([[xr, yr]]))
image_row = image_row + (x_reconstructed[0].reshape(28, 28), )
image_rows = image_rows + (np.hstack(image_row), )
image = np.vstack(image_rows)
plt.clf()
plt.imshow(image, cmap='Greys')
plt.tick_params(axis='both',
which='both',
bottom='off',
top='off',
labelbottom='off',
right='off',
left='off',
labelleft='off')
plt.show()
#Load our dataset
train_dataset = dutil.byPath(
args.dataset) if args.dataset else dutil.MNIST(False)
#Set up our data and parameters
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=1,
shuffle=True,
**kwargs)
num_items = train_dataset.data.size(0)
image_height = train_dataset.data.size(1)
image_width = train_dataset.data.size(2)
latent_features = 20
#Initialize the VAE
model = VAE(latent_features, image_height, image_width).to(device)
model.load_state_dict(torch.load(f'{MODEL_STORE_PATH}/{MODEL_NAME}'))
model.eval()
randoms = torch.FloatTensor(np.random.normal(0, 1, latent_features)).to(device)
some_image = model.decode(randoms)
plt.figure(1)
plt.title("Generated Image")
image_vector = some_image.view(-1, 28, 28).permute(1, 2, 0).cpu().squeeze()
fig2 = plt.imshow(image_vector.detach().numpy())
plt.show()
print single.color_scheme
latent_dim = 2
vae = VAE(784, [500], latent_dim)
vae.train(X,
batch_size=20,
num_epochs=100,
rerun=False,
model_filename='composite_l_%d' % latent_dim)
#xr = float(sys.argv[1])
#yr = float(sys.argv[2])
a = float(sys.argv[1])
b = float(sys.argv[2])
#c = float(sys.argv[3])
#d = float(sys.argv[4])
#e = float(sys.argv[5])
#x_reconstructed = vae.decode(np.array([[a,b,c,d,e]]))
x_reconstructed = vae.decode(np.array([[a, b]]))
plt.imshow(x_reconstructed[0].reshape(28, 28), cmap='Greys')
plt.tick_params(axis='both',
which='both',
bottom='off',
top='off',
labelbottom='off',
right='off',
left='off',
labelleft='off')
plt.box(on=False)
plt.savefig('../final_figures/%s' % sys.argv[3], bbox_inches='tight')
def load_model(model_filename, s0, s1):
"""Load a model which has already been trained.
NB we will be calling .predict() on the model. Keras will set
learning_phase correctly to test mode, as stated here: https://github.com/tensorflow/tensorflow/issues/11336#issuecomment-315898065
This is important because we have Dropout. """
if model_filename.endswith(".pt"):
# pytorch
import torch
from vae import VAE
device = torch.device("cpu")
model = VAE(s0, s1, 10, 400, 1.0, "BCE", "ReLU").to(device)
model.eval() # tell the model we are not in training mode
model.load_state_dict(torch.load(model_filename))
decode = lambda z: model.decode(torch.tensor(z, dtype=torch.float).flatten()).detach().numpy()
encode = lambda x: model.encodedecode(torch.tensor(x, dtype=torch.float).flatten()).detach().numpy()
recon = lambda x: model.forward(torch.tensor(x, dtype=torch.float).flatten()).detach().numpy() # TODO check this
latent_dim = 10
elif model_filename.endswith(".h5"):
import vae_conv_keras_drums
config_filename = model_filename.replace(".h5", "_args.txt")
s = open(config_filename).read()
args = eval(s)
# latent_dim, loss, Lambda, dropout_rate, conv_layers, filter_expansion, kernel_size0, kernel_size1 = config.split("_")
# config = model_filename.split("/")[-2]
# @dataclass
# class Namespace:
# Lambda: float=0.01
# batch_size: int=128
# conv_layers: int=2
# dense_layer_size: int=16
# dropout_rate: float=0.2
# epochs: int=100
# filter_expansion: int=2
# filters: int=16
# kernel_size0: int=9
# kernel_size1: int=4
# latent_dim: int=10
# loss: str='BCE'
# stride0: int=1
# stride1: int=1
# weights: np.array=np.array([[]])
# args = Args()
# args.latent_dim = int(latent_dim)
# args.loss = loss
# args.Lambda = float(Lambda)
# args.dropout_rate = float(dropout_rate)
# args.conv_layers = int(conv_layers)
# args.filter_expansion = int(filter_expansion)
# args.kernel_size0 = int(kernel_size0)
# args.kernel_size1 = int(kernel_size1)
# args.stride0 = 1
# args.stride1 = 1
# args.filters = 16
# args.dense_layer_size = 16
# args.epochs = 100
# args.batch_size = 128
input_shape = (9, 64, 1) # FIXME
inputs, outputs, z_mean, z_log_var, encoder, decoder, vae = vae_conv_keras_drums.make_model(input_shape, args)
vae.load_weights(model_filename)
encode = encoder.predict
decode = decoder.predict
recon = vae.predict
else:
raise
return encode, decode, recon, int(args.latent_dim)
def main():
##############
### PARAMS ###
##############
batch_train = 16
batch_test = 16
lr = 1e-4
beta1 = 0.9
epochs = 200
latent_dim = 200
nf = 64
model_name = 'vae-64fc-200nz-deeper'
val_every = 5
load_path = './saved_models/best_acc-vae-64fc-200nz-deeper-cosine-choose_dict.pt'
start_epoch = 1
transform = T.Compose([T.Resize(size=256),
T.CenterCrop(size=224),
T.ToTensor()
])
sample_dir = os.path.join('./results/samples/', model_name)
recon_dir = os.path.join('./results/reconstructions', model_name)
validate_params = {
'per_row': 5,
'metric': 'cosine',
'model': None,
'drop_last': None,
'root_dir': '/media/STORAGE/DATASETS/open-images/train/',
'csv_test': './csv_files/handpicked.csv',
'csv_val': './csv_files/validation_labeled.csv',
'batch_test': 1,
'batch_val': 1,
'save_test': True,
'save_val': True,
'vis_val': None,
'vis_test': None,
'algorithm': 'choose_dict'
}
##############################
### DIRS FOR SAVING IMAGES ###
##############################
try:
os.makedirs(sample_dir)
except FileExistsError:
print('{} already exists!'.format(sample_dir))
try:
os.makedirs(recon_dir)
except FileExistsError:
print('{} already exists!'.format(recon_dir))
##############################
### DATASETS & DATALOADERS ###
##############################
train_dataset_params = {
'root': '/media/STORAGE/DATASETS/open-images/train/',
'transform': transform,
}
test_dataset_params = {
'csv_file': './csv_files/test.csv',
'root_dir': '/media/STORAGE/DATASETS/open-images/challenge2018/',
'transform': transform,
}
train_dataloader_params = {
'batch_size': batch_train,
'shuffle': True,
'num_workers': 1
}
test_dataloader_params = {
'batch_size': batch_test,
'shuffle': True,
'num_workers': 1
}
train_dataset = ImageFolder(**train_dataset_params)
test_dataset = TestDataset(**test_dataset_params)
train_dataloader = DataLoader(train_dataset, **train_dataloader_params)
test_dataloader = DataLoader(test_dataset, **test_dataloader_params)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = VAE(latent_dim=latent_dim, nf=nf).to(device)
optimizer = optim.Adam(model.parameters(), lr=lr, betas=(beta1, 0.999))
best_train_loss = 99999999
best_acc = 0
if load_path:
checkpoint = torch.load(load_path)
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
epoch = checkpoint['epoch']
start_epoch = epoch + 1
best_train_loss = checkpoint['train_loss']
best_acc = checkpoint['acc_val']
print('Loaded model with train loss: {}, val_acc: {}'.format(checkpoint['train_loss'], checkpoint['acc_val']))
#####################
### TRAINING LOOP ###
#####################
for epoch in range(start_epoch, epochs + start_epoch):
train_loss = train(model, train_dataloader, optimizer, epoch, device, recon_dir)
test_loss = test(model, test_dataloader, epoch, device, recon_dir)
# train_loss = 10000
# 64 sets of random latent_dim-float vectors, i.e 64 locations
# in latent space
sample = torch.randn(36, latent_dim).to(device)
sample = model.decode(sample).detach().cpu()
# save out as an 8x8 matrix
# this will give you a visual idea of how well latent space can generate things
save_image(sample, os.path.join(sample_dir, 'sample_' + str(epoch) + '.png'))
###########################
### CHARADES VALIDATION ###
###########################
if epoch % val_every == 0:
alg = validate_params['algorithm']
metric = validate_params['metric']
validate_params['model'] = model.get_feature_extractor()
validate_params['vis_val'] = './vis_val/' + model_name + '/epoch' + str(epoch) + '/' + metric + '-' + alg
validate_params['vis_test'] = './vis_test/' + model_name + '/epoch' + str(epoch) + '/' + metric + '-' + alg
acc_val, acc_test = validate(**validate_params)
print(50 * '*')
print('CHARADES ===> Val acc: {}, Test acc: {}'.format(acc_val, acc_test))
print(50 * '*')
if best_acc < acc_val:
PATH = './saved_models/' + 'best_acc-' + model_name + '-' + metric + '-' + alg + '.pt'
torch.save({
'epoch': epoch,
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
'train_loss': best_train_loss,
'test_loss': test_loss,
'acc_val': best_acc,
'acc_test': acc_test
}, PATH)
print('Saving best acc model!')
if train_loss < best_train_loss:
alg = validate_params['algorithm']
metric = validate_params['metric']
PATH = './saved_models/' + 'best_train-' + model_name + metric + '-' + alg + '.pt'
torch.save({
'epoch': epoch,
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
'train_loss': best_train_loss,
'test_loss': test_loss,
'acc_val': best_acc,
'acc_test': -1
}, PATH)
print('Saving best train loss model!')
class BSSRDF:
"""BSSRDF class with VAE"""
def __init__(self, model_name):
"""Instanciation the VAE class and load a trained model"""
self.config = vae_config.VAEConfiguration()
self.device = torch.device("cuda")
# Instanciate and load trained model
model_path = f"{self.config.MODEL_DIR}\\{model_name}.pt"
self.model = VAE(self.config).to(self.device)
print(
f"model[{model_name}] is specified ({os.path.exists(model_path)})")
if os.path.exists(model_path):
self.model.load_state_dict(torch.load(model_path))
print(f"model[{model_name}] is loaded")
def estimate(self, in_pos, im, props, sigma_n, active):
"""
Estimate output position and absorption with VAE
Notice that the types of arguments are in pytorch (tensor) except sigma_n, but the ones of returns are in mitsuba (Vector3f, Float)
Args:
in_pos: Incident position in local mesh coordinates, and the type of this is tensor
im: Height map around incident position ranging to multiple of sigma_n.
This map can be generated by clip_scaled_map() from data_handler, and the type is tensor
props: Medium properties vector including (and following this order)
- effective albedo
- g
- eta
- incident angle (xyz)
- max height
, and the type of this argument is tensor
sigma_n: Standard deviation of the range of medium scattering.
In this method, this value is used as scale factor of coordinates in vae
active: Boolean mask which indicates whether a ray is applied VAE or not
Return:
recon_pos: estimated outgoing position (Vector3f)
recon_abs: estimated absorption probability (Spectrum)
"""
n_sample, _, _, _ = im.shape
pos = Vector3f(in_pos)
abs_prob = Float().zero(n_sample)
self.model.eval()
with torch.no_grad():
# Feature conversion
feature = self.model.feature_conversion(
im.to(self.device, dtype=torch.float),
props.to(self.device, dtype=torch.float))
# Sample latent variable from normal distribution
z = torch.randn(n_sample, 4).to(self.device, dtype=torch.float)
# Decode and get reconstructed position and absorption
recon_pos, recon_abs = self.model.decode(feature, z)
# Convert from tensor to Vector3f and Float
recon_pos = Vector3f(recon_pos)
abs_prob = Spectrum(recon_abs.view(1, -1).squeeze())
# Reconstruct real scale position in mesh local coordinates
pos += ek.select(active, sigma_n * recon_pos, 0)
return pos, abs_prob
def sample_bssrdf(self, scene, bsdf, bs, si, bdata, heightmap_pybind,
channel, active):
"""
Get projected sample position and absorption probability form VAE BSSRDF
Args:
scene: rendered scene object
bsdf: BSDF information of each ray
bs: BSDF Sample object of each ray
si: Surface interaction object of each ray
bdata: BSSRDF data object in scene data from data_pipeline.py
channel: RGB channel of interest
TODO: Support multi channel sampling
active: Mask data whether a ray needs to process BSSRDF or not
Returns:
projected_si: Projected surface interaction object from BSSRDF
proj_suc: Mask data indicating projection succeed or not
abs_recon: absorption probability from BSSRDF
"""
# Get ID of BSSDRF mesh for each sur face interactions
mesh_id = BSDF.mesh_id_vec(bsdf, active)
# Convert incident position into local coordinates of mesh of interested as tensor
in_pos = ek.select(active, si.to_mesh_local(bs), Vector3f(0))
# Get properties, e.g., medium params and incident angle as tensor
props, sigma_n = get_props(bs, si, channel)
# Get height map around incident position as tensor
im_bind = heightmap_pybind.get_height_map(in_pos.torch().cpu(),
mesh_id.torch().cpu())
im = torch.tensor(im_bind)
# Estimate position and absorption probability with VAE as mitsuba types
recon_pos_local, abs_recon = self.estimate(in_pos.torch(), im, props,
sigma_n, active)
# Convert from mesh coordinates to world coordinates
recon_pos_world = si.to_mesh_world(bs, recon_pos_local)
# Project estimated position onto nearest mesh
projected_si, proj_suc = si.project_to_mesh_normal(
scene, recon_pos_world, bs, channel, active)
return projected_si, proj_suc, abs_recon
def test(epoch):
model.eval()
losses = []
with torch.no_grad():
for i, (data, _) in enumerate(test_loader):
data = data.to(device)
recon_batch, mu, logvar = model(data)
bce, kld = loss_function(recon_batch, data, mu, logvar)
loss = bce+kld
losses.append(loss.item())
if i == 0:
n = min(data.size(0), 8)
comparison = torch.cat([data[:n],
recon_batch.view(opts.batchSize, 1, 28, 28)[:n]])
save_image(comparison.cpu(),
os.path.join(model_path, 'reconstruction_' + str(epoch) + '.png'), nrow=n)
writer.add_scalar('nll_epoch', np.array(losses).mean(), epoch)
if __name__ == "__main__":
from scipy.stats import norm
grid_x = norm.ppf(np.linspace(0.0000001, 0.9999999, 64))
sample = torch.tensor(grid_x).float().to(device).view(64,1)
# sample = torch.randn(64, opts.nz).to(device)
for epoch in range(0, opts.epochs):
train(epoch)
test(epoch)
with torch.no_grad():
image = model.decode(sample).cpu()
image = make_grid(image.view(64,1,28,28), nrow=8, normalize=True, scale_each=True)
writer.add_image('random images', image, epoch)
class Interface:
def __init__(self):
self.n_latent = 16
self.update = True
self.root = tk.Tk()
self.root.geometry("800x1200")
self.model = VAE(n_latent=16)
self.model.load()
self.latent_activations = np.zeros((1, 16))
img = self.model.decode(self.latent_activations)
self.imfig = plt.figure()
self.imax = self.imfig.add_subplot(111)
self.imax.set_axis_off()
self.imax.imshow(img.numpy().squeeze())
canvas = FigureCanvasTkAgg(self.imfig, master=self.root)
canvas.draw()
canvas.get_tk_widget().pack()
btn = tk.Button(master=self.root,
text="randomize",
command=self.randomize)
btn.pack()
self.sliders = list()
for i in range(self.n_latent):
slider = tk.Scale(master=self.root,
command=self.update_latent,
from_=-2,
to=2,
orient=tk.HORIZONTAL,
resolution=0.01,
length=500)
slider.pack()
self.sliders.append(slider)
self.randomize()
self.root.mainloop()
def update_image(self):
if self.update:
img = self.model.decode(self.latent_activations)
img = img.numpy().squeeze()
self.imax.imshow(img)
self.imfig.canvas.draw()
def randomize(self):
self.latent_activations = np.random.random((1, 16))
self.update = False
for i in range(self.n_latent):
self.sliders[i].set(self.latent_activations[0, i])
self.update = True
self.update_image()
def update_latent(self, _):
if self.update:
for i in range(self.n_latent):
self.latent_activations[0, i] = self.sliders[i].get()
self.update_image()
def train(save_file, dataset):
# Device configuration
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# Create a directory if not exists
sample_dir = 'samples'
if not os.path.exists(sample_dir):
os.makedirs(sample_dir)
# Hyper-parameters
image_size = 32 * 32 * 3
z_dim = 512
num_epochs = 100
batch_size = 128
learning_rate = 1e-3
# Data loader
data_loader = torch.utils.data.DataLoader(dataset=dataset,
batch_size=batch_size,
shuffle=True)
model = VAE().to(device)
model.train()
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.85)
# Start training
for epoch in range(num_epochs):
scheduler.step()
for i, (x, _) in enumerate(data_loader):
# Forward pass
x = x.to(device) #.view(-1, image_size)
x_reconst, mu, log_var = model(x)
# Compute reconstruction loss and kl divergence
# For KL divergence, see Appendix B in VAE paper or http://yunjey47.tistory.com/43
reconst_loss = F.mse_loss(x_reconst, x, size_average=False)
kl_div = -0.5 * torch.sum(1 + log_var - mu.pow(2) - log_var.exp())
# Backprop and optimize
loss = reconst_loss #+ kl_div
optimizer.zero_grad()
loss.backward()
optimizer.step()
if (i + 1) % 10 == 0:
print(
"Epoch[{}/{}], Step [{}/{}], Reconst Loss: {:.4f}, KL Div: {:.4f}"
.format(epoch + 1, num_epochs, i + 1, len(data_loader),
reconst_loss.item(), kl_div.item()))
with torch.no_grad():
# Save the sampled images
z = torch.randn(batch_size, z_dim).to(device)
out = model.decode(z).view(-1, 3, 32, 32)
save_image(
out,
os.path.join(sample_dir, 'sampled-{}.png'.format(epoch + 1)))
# Save the reconstructed images
out, _, _ = model(x)
x_concat = torch.cat(
[x.view(-1, 3, 32, 32),
out.view(-1, 3, 32, 32)], dim=3)
save_image(
x_concat,
os.path.join(sample_dir, 'reconst-{}.png'.format(epoch + 1)))
# Save model
torch.save(model, save_file)
def main():
# check if cuda available
device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
# define dataset and dataloader
dataset = AirfoilDataset()
airfoil_x = dataset.get_x()
airfoil_dim = airfoil_x.shape[0]
airfoil_dataloader = DataLoader(dataset, batch_size=16, shuffle=True)
# hyperparameters
latent_dim = 16 # please do not change latent dimension
lr = 0.001 # learning rate
num_epochs = 30
# build the model
vae = VAE(airfoil_dim=airfoil_dim, latent_dim=latent_dim).to(device)
print("VAE model:\n", vae)
# define optimizer for discriminator and generator separately
optim = Adam(vae.parameters(), lr=lr)
lossess = []
# train the VAE model
for epoch in range(num_epochs):
losses = []
for n_batch, (local_batch, __) in enumerate(airfoil_dataloader):
y_real = local_batch.to(device)
# train VAE
genData, M, logvar = vae(y_real)
loss = vae.computeLoss(M,logvar,genData,y_real)
losses.append(loss.item())
optim.zero_grad()
loss.backward()
optim.step()
# print loss while training
if (n_batch + 1) % 30 == 0:
print("Epoch: [{}/{}], Batch: {}, loss: {}".format(
epoch, num_epochs, n_batch, loss.item()))
lossess.append(np.mean(np.array(losses)))
plt.title('VAE loss')
plt.ylabel('Loss')
plt.xlabel('Num of Epochs')
plt.plot(lossess)
# test trained VAE model
num_samples = 100
# reconstuct airfoils
real_airfoils = dataset.get_y()[:num_samples]
recon_airfoils, __, __ = vae(torch.from_numpy(real_airfoils).to(device))
if 'cuda' in device:
recon_airfoils = recon_airfoils.detach().cpu().numpy()
else:
recon_airfoils = recon_airfoils.detach().numpy()
# randomly synthesize airfoils
noise = torch.randn((num_samples, latent_dim)).to(device) # create random noise
gen_airfoils = vae.decode(noise)
if 'cuda' in device:
gen_airfoils = gen_airfoils.detach().cpu().numpy()
else:
gen_airfoils = gen_airfoils.detach().numpy()
# plot real/reconstructed/synthesized airfoils
plot_airfoils(airfoil_x, real_airfoils)
plot_airfoils(airfoil_x, recon_airfoils)
plot_airfoils(airfoil_x, gen_airfoils)
import numpy as np
import input_quickdraw
import matplotlib.pyplot as plt
from vae import VAE
from scipy.stats import kde
single = input_quickdraw.Single(['alarm_clock'], 5000)
X = single.train_data
Y = single.train_label
vae = VAE(784, [500], 10)
vae.train(X, batch_size=200, num_epochs=100, rerun=False, model_filename='alarm_clock_schematic')
print('Encoding...')
X = X[:1, :]
plt.imshow(X[0].reshape(28,28), cmap='Greys')
plt.tick_params(axis='both', which='both', bottom='off', top='off', labelbottom='off', right='off', left='off', labelleft='off')
plt.show()
latent_z = vae.encode(X)
x_reconstructed = vae.decode(latent_z)
plt.imshow(x_reconstructed[0].reshape(28,28), cmap='Greys')
plt.tick_params(axis='both', which='both', bottom='off', top='off', labelbottom='off', right='off', left='off', labelleft='off')
plt.show()
def train_from_folder(data_dir="../data/",
results_dir="../data/results/",
models_dir="../data/models",
image_size=28,
batch_size=256,
num_epochs=15,
learning_rate=1e-3,
hidden_dim=400,
latent_dim=20,
save_every=10,
seed=42,
amp=False,
device="cuda"):
data_dir = Path(data_dir)
sample_dir = Path(results_dir)
models_dir = Path(models_dir)
sample_dir.mkdir(exist_ok=True)
models_dir.mkdir(exist_ok=True)
device = torch.device(device)
dataset = datasets.MNIST(root=data_dir,
train=True,
transform=transforms.ToTensor(),
download=True)
dataloader = data.DataLoader(dataset=dataset,
batch_size=batch_size,
shuffle=True)
model = VAE(image_size**2, hidden_dim, latent_dim).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
for epoch in range(num_epochs):
for i, (x, _) in enumerate(dataloader):
x = x.to(device).view(-1, image_size**2)
x_reconstructed, mu, log_var = model(x)
reconst_loss = F.binary_cross_entropy(x_reconstructed,
x,
size_average=False)
kl_div = -0.5 * torch.sum(1 + log_var - mu.pow(2) - log_var.exp())
total_loss = reconst_loss + kl_div
optimizer.zero_grad()
total_loss.backward()
optimizer.step()
if (i + 1) % 10 == 0:
print(
f"Epoch [{epoch + 1}/{num_epochs}], Steps [{i + 1}/{len(dataloader)}], "
f"Reconstruction loss: {reconst_loss.item():.4f}, KL div: {kl_div.item():.4f}"
)
with torch.no_grad():
z = torch.randn(batch_size, latent_dim).to(device)
out = model.decode(z).view(-1, 1, 28, 28)
save_image(out, Path(sample_dir) / f"sampled-{epoch + 1}.png")
out, _, _ = model(x)
x_concat = torch.cat(
[x.view(-1, 1, 28, 28),
out.view(-1, 1, 28, 28)], dim=3)
save_image(x_concat, Path(sample_dir) / f"reconst-{epoch + 1}.png")
def train(config):
device = torch.device(
'cuda') if torch.cuda.is_available() else torch.device('cpu')
train_dataset = create_mnist_dataset(config.dataset_path,
config.batch_size, True)
eval_dataset = create_mnist_dataset(config.dataset_path, config.batch_size,
False)
net = VAE().to(device)
full_model_path = config.model_path + config.model_name
if config.load_pretrained_model:
net.load_state_dict(torch.load(full_model_path))
print('Load the pretrained model from %s successfully!' %
full_model_path)
else:
weight_init(net)
if not os.path.exists(config.model_path):
os.makedirs(config.model_path)
print('First time training!')
net.train()
optimizer = torch.optim.Adam(net.parameters(),
lr=config.learning_rate,
betas=[0.5, 0.999])
summary = SummaryWriter(config.summary_path)
total_iter = 1
for e in range(1, config.epoch + 1):
for idx, (x, _) in enumerate(train_dataset):
x = x.to(device).view(-1, 784)
optimizer.zero_grad()
recon_x, mu, logvar = net(x)
loss = loss_func(recon_x, x, mu, logvar)
loss.backward()
optimizer.step()
print('[Epoch %d|Train Batch %d] Loss = %.6f' %
(e, idx, loss.item()))
summary.add_scalar('Train/Loss', loss.item(), total_iter)
total_iter += 1
if e % 5 == 0:
net.eval()
eval_losses = []
with torch.no_grad():
for idx, (x, _) in enumerate(eval_dataset):
x = x.to(device).view(-1, 784)
recon_x, mu, logvar = net(x)
loss = loss_func(recon_x, x, mu, logvar)
print('[Epoch %d|Eval Batch %d] Loss = %.6f' %
(e, idx, loss.item()))
eval_losses.append(loss.item())
mean_eval_loss = np.mean(eval_losses)
summary.add_scalar('Eval/Loss', mean_eval_loss, e)
net.train()
if e % 5 == 0:
with torch.no_grad():
fake_z = torch.randn((64, net.nz)).to(device)
fake_imgs = net.decode(fake_z).view(-1, 1, 28, 28).detach()
fake_imgs = torchvision.utils.make_grid(
fake_imgs, padding=2,
normalize=True).detach().cpu().numpy()
summary.add_image('Eval/Fake_imgs_after_%d_epochs' % e,
fake_imgs, e)
if e % 2 == 0:
torch.save(net.state_dict(), full_model_path)
summary.close()
