def __init__(self, disc_hiddens=[1000, 1000, 1000],
gamma=30, input_size=(3, 64, 64),
kernel_sizes=[32, 32, 64, 64],
hidden_size=256, dim_z=32,
binary=True, **kwargs):
"""initialize neural networks
:param disc_hiddens: list of int, numbers of hidden units of each layer in discriminator
:param gamma: weight for total correlation term in loss function
"""
super(ConvFactorVAE, self).__init__()
self.gamma = gamma
self.dim_z = dim_z
self.binary = binary
self.input_size = input_size
self.hidden_size = hidden_size
# VAE networks
self.vae = ConvVAE(input_size, kernel_sizes, hidden_size, dim_z, binary, **kwargs)
# inherit some attributes
self.channel_sizes = self.vae.channel_sizes
# discriminator networks
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, disc_hiddens, act_layer=D_act, act_args=D_act_args)
self.discriminator = nn.Sequential(
self.discriminator,
nn.Linear(disc_hiddens[-1], D_output_dim))
def __init__(self, load_model=True):
self.env_name = "carracing"
self.vae = ConvVAE(batch_size=1,
gpu_mode=False,
is_training=False,
reuse=True)
self.rnn = MDNRNN(hps_sample, gpu_mode=False, reuse=True)
if load_model:
self.vae.load_json('Weights/vae_weights.json')
self.rnn.load_json('Weights/rnn_weights.json')
self.state = rnn_init_state(self.rnn)
self.rnn_mode = True
self.input_size = rnn_output_size(EXP_MODE)
self.z_size = 32
if EXP_MODE == MODE_Z_HIDDEN:
self.hidden_size = 40
self.weight_hidden = np.random.randn(self.input_size,
self.hidden_size)
self.bias_hidden = np.random.randn(self.hidden_size)
self.weight_output = np.random.randn(self.hidden_size, 3)
self.bias_output = np.random.randn(3)
self.param_count = ((self.input_size + 1) *
self.hidden_size) + (self.hidden_size * 3 + 3)
else:
self.weight = np.random.randn(self.input_size, 3)
self.bias = np.random.randn(3)
self.param_count = (self.input_size) * 3 + 3
self.render_mode = False
kl_tolerance = 0.5
filelist = os.listdir(DATA_DIR)
filelist.sort()
filelist = filelist[0:10000]
dataset, action_dataset = load_raw_data_list(filelist)
# Resetting the graph of the VAE model
reset_graph()
# Creating the VAE model as an object of the ConvVAE class
vae = ConvVAE(z_size=z_size,
batch_size=batch_size,
learning_rate=learning_rate,
kl_tolerance=kl_tolerance,
is_training=False,
reuse=False,
gpu_mode=True)
# Loading the weights of the VAE model
vae.load_json(os.path.join(model_path_name, 'vae.json'))
# Running the main code that generates the data from the VAE model for the MDN-RNN model
mu_dataset = []
logvar_dataset = []
for i in range(len(dataset)):
data_batch = dataset[i]
mu, logvar, z = encode_batch(data_batch)
os.makedirs(out_path, exist_ok=True)
# check for GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# load data
print("Loading data")
train_dataloader = get_dataloader(opt)
test_dataloader = None
n = len(train_dataloader.dataset)
iter_per_epoch = math.ceil(n / opt.batch_size)
# run
start = time.time()
print("Training")
vae = ConvVAE(opt).to(device)
optimizer = optim.Adam(vae.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))
discriminator = Discriminator(opt.latent_dim).to(device)
optimizer_d = optim.Adam(discriminator.parameters(),
lr=opt.lrd,
betas=(opt.b1d, opt.b2d))
losses_train = []
reconstruction_error = []
kl_divergence = []
total_correlation = []
discriminator_loss = []
losses_test = []
dimension_kld = run()
def main():
parser = argparse.ArgumentParser(description='VAE')
parser.add_argument(
'--weights',
type=str,
default=None,
help='Load weigths for the VAE, not training the model')
parser.add_argument('--batch-size',
type=int,
default=128,
help='Input batch size for training, 128 by default')
parser.add_argument('--epochs',
type=int,
default=10,
help='Number of epochs to train, 10 by default')
parser.add_argument('--seed',
type=int,
default=1,
help='Random seed, 1 by default')
parser.add_argument(
'--log-interval',
type=int,
default=10,
help='Interval between two logs of the training status')
args = parser.parse_args()
torch.manual_seed(args.seed)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
if torch.cuda.is_available():
kwargs = {'num_workers': 1, 'pin_memory': True}
else:
kwargs = {}
train_dataset = datasets.MNIST('data',
train=True,
download=True,
transform=transforms.ToTensor())
test_dataset = datasets.MNIST('data',
train=False,
transform=transforms.ToTensor())
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=args.batch_size,
shuffle=True,
**kwargs)
test_loader = torch.utils.data.DataLoader(test_dataset,
batch_size=args.batch_size,
shuffle=True,
**kwargs)
# model = FCVAE(channels=1, z_dim=2).to(device)
model = ConvVAE(channels=1, z_dim=2).to(device)
net_manager = NetManager(model,
device,
train_loader=train_loader,
test_loader=test_loader)
if args.weights is None:
net_manager.set_writer("vae")
net_manager.train(args.epochs, log_interval=args.log_interval)
if not os.path.exists("weights"):
os.mkdir("weights")
net_manager.save_net("weights/vae.pth")
else:
net_manager.load_net(args.weights)
net_manager.plot_latent_space(dark_background=True)
crops = np.load(args["npy"])
# parameters
z_size=16
batch_size, crop_size, _, _ = crops.shape
print("batch_size {}".format(batch_size))
print("crop_size {}".format(crop_size))
cv2.imshow("crop",crops[0])
cv2.waitKey(0)
reset_graph()
test_vae = ConvVAE(z_size=z_size,
batch_size=batch_size,
is_training=False,
reuse=False,
gpu_mode=True)
# show reconstruction example
test_vae.load_json("../../models/0/vae_{}.json".format(180000))
z = test_vae.encode(crops)
print(z.shape)
rec = test_vae.decode(z)
print(rec.shape)
np.save("../../output/z_{}.npy".format(batch_size), z)
np.save("../../output/rec_{}.npy".format(batch_size), rec)
#for img_idx, img in enumerate(test_batch):
# run this file from the script folder
data_path = "../data/dsprites/dsprites_ndarray_co1sh3sc6or40x32y32_64x64.npz"
n_train = 10000
n_test = 5000
path = "../results/dsprites/"
metrics = {"train": None, "test": None}
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model_path = path + args.model + f"/parameter_{args.parameter}/seed_{args.seed}"
model_data = torch.load(model_path + "/model.pt",
map_location=torch.device("cpu"))
model = ConvVAE(model_data["opt"]).to(device)
model.load_state_dict(state_dict=model_data["model"])
opt = model_data["opt"]
evaluate = Evaluator(data_path=data_path,
model=model,
opt=opt,
batch_size=64)
if args.metric == "beta":
metric_train, metric_test = evaluate.get_beta_metric(n_train, n_test)
metrics["train"] = metric_train
metrics["test"] = metric_test
elif args.metric == "factor":
metric_train, metric_test = evaluate.get_factor_metric(