Spaces:
Running
Running
| #!/usr/bin/env python3 | |
| """Code to generate plots for Extended Data Fig. 4.""" | |
| import os | |
| import matplotlib | |
| import matplotlib.pyplot as plt | |
| import numpy as np | |
| import echonet | |
| def main(root=os.path.join("timing", "video"), | |
| fig_root=os.path.join("figure", "complexity"), | |
| FRAMES=(1, 8, 16, 32, 64, 96), | |
| pretrained=True): | |
| """Generate plots for Extended Data Fig. 4.""" | |
| echonet.utils.latexify() | |
| os.makedirs(fig_root, exist_ok=True) | |
| fig = plt.figure(figsize=(6.50, 2.50)) | |
| gs = matplotlib.gridspec.GridSpec(1, 3, width_ratios=[2.5, 2.5, 1.50]) | |
| ax = (plt.subplot(gs[0]), plt.subplot(gs[1]), plt.subplot(gs[2])) | |
| # Create legend | |
| for (model, color) in zip(["EchoNet-Dynamic (EF)", "R3D", "MC3"], matplotlib.colors.TABLEAU_COLORS): | |
| ax[2].plot([float("nan")], [float("nan")], "-", color=color, label=model) | |
| ax[2].set_title("") | |
| ax[2].axis("off") | |
| ax[2].legend(loc="center") | |
| for (model, color) in zip(["r2plus1d_18", "r3d_18", "mc3_18"], matplotlib.colors.TABLEAU_COLORS): | |
| for split in ["val"]: # ["val", "train"]: | |
| print(model, split) | |
| data = [load(root, model, frames, 1, pretrained, split) for frames in FRAMES] | |
| time = np.array(list(map(lambda x: x[0], data))) | |
| n = np.array(list(map(lambda x: x[1], data))) | |
| mem_allocated = np.array(list(map(lambda x: x[2], data))) | |
| # mem_cached = np.array(list(map(lambda x: x[3], data))) | |
| batch_size = np.array(list(map(lambda x: x[4], data))) | |
| # Plot Time (panel a) | |
| ax[0].plot(FRAMES, time / n, "-" if pretrained else "--", marker=".", color=color, linewidth=(1 if split == "train" else None)) | |
| print("Time:\n" + "\n".join(map(lambda x: "{:8d}: {:f}".format(*x), zip(FRAMES, time / n)))) | |
| # Plot Memory (panel b) | |
| ax[1].plot(FRAMES, mem_allocated / batch_size / 1e9, "-" if pretrained else "--", marker=".", color=color, linewidth=(1 if split == "train" else None)) | |
| print("Memory:\n" + "\n".join(map(lambda x: "{:8d}: {:f}".format(*x), zip(FRAMES, mem_allocated / batch_size / 1e9)))) | |
| print() | |
| # Labels for panel a | |
| ax[0].set_xticks(FRAMES) | |
| ax[0].text(-0.05, 1.10, "(a)", transform=ax[0].transAxes) | |
| ax[0].set_xlabel("Clip length (frames)") | |
| ax[0].set_ylabel("Time Per Clip (seconds)") | |
| # Labels for panel b | |
| ax[1].set_xticks(FRAMES) | |
| ax[1].text(-0.05, 1.10, "(b)", transform=ax[1].transAxes) | |
| ax[1].set_xlabel("Clip length (frames)") | |
| ax[1].set_ylabel("Memory Per Clip (GB)") | |
| # Save figure | |
| plt.tight_layout() | |
| plt.savefig(os.path.join(fig_root, "complexity.pdf")) | |
| plt.savefig(os.path.join(fig_root, "complexity.eps")) | |
| plt.close(fig) | |
| def load(root, model, frames, period, pretrained, split): | |
| """Loads runtime and memory usage for specified hyperparameter choice.""" | |
| with open(os.path.join(root, "{}_{}_{}_{}".format(model, frames, period, "pretrained" if pretrained else "random"), "log.csv"), "r") as f: | |
| for line in f: | |
| line = line.split(",") | |
| if len(line) < 4: | |
| # Skip lines that are not csv (these lines log information) | |
| continue | |
| if line[1] == split: | |
| *_, time, n, mem_allocated, mem_cached, batch_size = line | |
| time = float(time) | |
| n = int(n) | |
| mem_allocated = int(mem_allocated) | |
| mem_cached = int(mem_cached) | |
| batch_size = int(batch_size) | |
| return time, n, mem_allocated, mem_cached, batch_size | |
| raise ValueError("File missing information.") | |
| if __name__ == "__main__": | |
| main() | |