File size: 2,663 Bytes
59630ba | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 | import random
from pathlib import Path
import torch
from tqdm import tqdm
import argparse
from algorithms.vae.common.distribution import DiagonalGaussianDistribution
def estimate_latent_stats(
latent_dir: str,
batch_size: int = 128,
repeat: int = 5,
channel_wise: bool = False,
is_distribution: bool = False,
):
data_dir = Path(latent_dir)
latent_paths = [
p for p in data_dir.glob("**/*.pt") if not p.name.endswith("_cond.pt")
]
random.shuffle(latent_paths)
means, stds = [], []
for i in range(repeat):
latents = None
for idx in tqdm(
range(batch_size * i, batch_size * (i + 1)), desc=f"Batch {i + 1}"
):
latent = torch.load(latent_paths[idx])
if is_distribution:
latent = DiagonalGaussianDistribution(latent).sample()
if latents is None:
latents = latent
else:
latents = torch.cat([latents, latent], dim=0)
if channel_wise:
mean = latents.mean(dim=(0, 2, 3)).tolist() # Compute mean for each channel
std = latents.std(dim=(0, 2, 3)).tolist() # Compute std for each channel
else:
mean = latents.mean().item()
std = latents.std().item()
means.append(mean)
stds.append(std)
print(f"Batch {i + 1}: mean: {mean}, std: {std}")
if channel_wise:
overall_mean = [
sum(m[i] for m in means) / len(means) for i in range(len(means[0]))
]
overall_std = [sum(s[i] for s in stds) / len(stds) for i in range(len(stds[0]))]
overall_mean = [round(x, 3) for x in overall_mean]
overall_std = [round(x, 3) for x in overall_std]
else:
overall_mean = sum(means) / len(means)
overall_std = sum(stds) / len(stds)
print(f"Overall mean: {overall_mean}, overall std: {overall_std}")
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--latent_dir", type=str, help="Path to the latent tensors")
parser.add_argument("--batch_size", type=int, default=128, help="Batch size")
parser.add_argument("--repeat", type=int, default=5, help="Number of repeats")
parser.add_argument(
"--channel_wise", action="store_true", help="Calculate statistics channel wise"
)
parser.add_argument(
"--distribution", action="store_true", help="Latents are saved as distributions"
)
args = parser.parse_args()
estimate_latent_stats(
args.latent_dir,
args.batch_size,
args.repeat,
args.channel_wise,
args.distribution,
)
|