def load_model( model_type, max_filters, num_layers, image_size, model_parameters, small_conv, model_path, device, ): if model_type == "vae": model = vae.ConvolutionalVAE( max_filters=max_filters, num_layers=num_layers, input_image_dimensions=image_size, latent_dim=model_parameters, small_conv=small_conv, ) elif model_type == "dual_input_vae": model = vae.FusionVAE( max_filters=max_filters, num_layers=num_layers, input_image_dimensions=image_size, latent_dim=model_parameters, small_conv=small_conv, ) elif model_type == "vq_vae": num_embeddings = model_parameters["K"] embedding_dim = model_parameters["D"] commitment_cost = model_parameters["commitment_cost"] model = vqvae.VQVAE( num_layers=num_layers, input_image_dimensions=image_size, small_conv=small_conv, embedding_dim=embedding_dim, num_embeddings=num_embeddings, commitment_cost=commitment_cost, ) elif model_type == "dual_input_autoencoder": model = autoencoder.FusionAE( max_filters=max_filters, num_layers=num_layers, input_image_dimensions=image_size, latent_dim=model_parameters, small_conv=small_conv, ) else: model = autoencoder.ConvolutionalAE( max_filters=max_filters, num_layers=num_layers, input_image_dimensions=image_size, latent_dim=model_parameters, small_conv=small_conv, ) model.load_state_dict(torch.load(model_path, map_location=device)) model.eval() return model
# Create & Load VQVAE Model embedding_model = vqvae.VQVAE( num_layers=vq_vae_num_layers, input_image_dimensions=image_size, small_conv=vq_vae_small_conv, embedding_dim=vq_vae_embedding_dim, num_embeddings=vq_vae_num_embeddings, commitment_cost=vq_vae_commitment_cost, use_max_filters=vq_vae_use_max_filters, max_filters=vq_vae_max_filters, ) else: embedding_model = vae.ConvolutionalVAE( max_filters=vae_max_filters, num_layers=vae_num_layers, input_image_dimensions=image_size, latent_dim=vae_latent_dim, small_conv=vae_small_conv, ) embedding_model.load_state_dict(torch.load(embedding_model_path, map_location=device)) embedding_model.eval() embedding_model.to(device) # Setup Label Encoder all_classes = train_data.get_classes() num_classes = len(all_classes) label_encoder = LabelEncoder().fit(all_classes) # Create Model model = classifier.ANNMultiClassClassifier( num_layers=num_layers,
if not os.path.exists(generated_dir): os.makedirs(generated_dir) # Setup Device gpu = torch.cuda.is_available() device = torch.device("cuda" if gpu else "cpu") print(gpu, device) # Preprocess & Create Data Loaders transform = data.image2tensor_resize(image_size) # Create Model model = vae.ConvolutionalVAE( max_filters=max_filters, num_layers=num_layers, input_image_dimensions=image_size, latent_dim=latent_dim, small_conv=small_conv, ) model.load_state_dict(torch.load(model_path, map_location=device)) model.eval() model.to(device) with torch.no_grad(): for iteration in tqdm(range(num_sample_batches)): embeddings = embeddings = torch.rand((batch_size, latent_dim), device=device) generated = model.decoder(embeddings) generated = generated.permute(0, 2, 3, 1).detach().cpu().numpy() for j, image in enumerate(generated):
num_embeddings=vq_vae_num_embeddings, commitment_cost=vq_vae_commitment_cost, use_max_filters=vq_vae_use_max_filters, max_filters=vq_vae_max_filters, ) vq_vae.load_state_dict(torch.load(vq_vae_model_path, map_location=device)) vq_vae.eval() vq_vae.to(device) # Create Model if mode == "continuous": if model_type == "vae": model = vae.ConvolutionalVAE( image_channels=input_channels, max_filters=max_filters, num_layers=num_layers, latent_dim=latent_dim, input_image_dimensions=input_dim, small_conv=False, ) else: model = cnn_prior.InfillingCNNPrior( num_layers=num_layers, max_filters=max_filters, input_channels=1 ) else: model = vae.VAEPrior( max_filters=max_filters, num_layers=num_layers, latent_dim=latent_dim, input_dimensions=input_dim,