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,