| 160 | plt.close() |
| 161 | |
| 162 | def plot_gradient_flow(named_parameters, save_path: Optional[str] = None): |
| 163 | ave_grads = [] |
| 164 | max_grads = [] |
| 165 | layers = [] |
| 166 | |
| 167 | for n, p in named_parameters: |
| 168 | if p.requires_grad and p.grad is not None: |
| 169 | layers.append(n) |
| 170 | ave_grads.append(p.grad.abs().mean().cpu().item()) |
| 171 | max_grads.append(p.grad.abs().max().cpu().item()) |
| 172 | |
| 173 | plt.figure(figsize=(12, 6)) |
| 174 | plt.bar(np.arange(len(max_grads)), max_grads, alpha=0.5, lw=1, color="c", label="max gradient") |
| 175 | plt.bar(np.arange(len(ave_grads)), ave_grads, alpha=0.5, lw=1, color="b", label="mean gradient") |
| 176 | plt.hlines(0, 0, len(ave_grads) + 1, lw=2, color="k") |
| 177 | plt.xticks(range(0, len(ave_grads), 1), layers, rotation="vertical") |
| 178 | plt.xlim(left=0, right=len(ave_grads)) |
| 179 | plt.ylim(bottom=-0.001, top=max(max_grads) * 1.1) |
| 180 | plt.xlabel("Layers") |
| 181 | plt.ylabel("Gradient") |
| 182 | plt.title("Gradient Flow") |
| 183 | plt.grid(True, alpha=0.3) |
| 184 | plt.legend() |
| 185 | plt.tight_layout() |
| 186 | |
| 187 | if save_path: |
| 188 | os.makedirs(os.path.dirname(save_path), exist_ok=True) |
| 189 | plt.savefig(save_path, dpi=300, bbox_inches='tight') |
| 190 | print(f"Gradient flow plot saved to {save_path}") |
| 191 | |
| 192 | plt.close() |