# Create Output Paths if not os.path.exists(output_dir): os.makedirs(output_dir) ################################################################################ ################################## Data Setup ################################## ################################################################################ # Load Data train = data.load_images_from_folder(train_data_folder, use_noise_images) val = data.load_images_from_folder(val_data_folder, use_noise_images) test = data.load_images_from_folder(test_data_folder, use_noise_images) # Preprocess & Create Data Loaders transform = data.image2tensor_resize(image_size) train_data = data.CustomDataset(train, transform) val_data = data.CustomDataset(val, transform) test_data = data.CustomDataset(test, transform) train_dataloader = torch.utils.data.DataLoader( train_data, batch_size=batch_size, shuffle=True, num_workers=num_dataloader_workers, pin_memory=gpu, ) val_dataloader = torch.utils.data.DataLoader( val_data, batch_size=batch_size,
gpu = torch.cuda.is_available() device = torch.device("cuda" if gpu else "cpu") print(gpu, device) # Create Output Paths if not os.path.exists(reconstructed_dir): os.makedirs(reconstructed_dir) if not os.path.exists(generated_dir): os.makedirs(generated_dir) ################################################################################ ################################## Data Setup ################################## ################################################################################ # Preprocess & Create Data Loaders transform = data.image2tensor_resize(vq_vae_image_size) # Load Data train_data = data.CustomDatasetNoMemory(train_data_folder, transform, use_noise_images) val_data = data.CustomDatasetNoMemory(val_data_folder, transform, use_noise_images) test_data = data.CustomDatasetNoMemory(test_data_folder, transform, use_noise_images) train_dataloader = torch.utils.data.DataLoader( train_data, batch_size=batch_size, shuffle=True, num_workers=num_dataloader_workers, pin_memory=gpu,