ptyrad.plotting.model

ptyrad.plotting.model#

Plotting functions related to PyTorch models

Functions

plot_forward_pass(model, indices, dp_power)

Plot the forward pass for the input torch model

plot_summary(output_path, model, niter, ...)

Wrapper function for most visualization function

ptyrad.plotting.model.plot_summary(output_path, model, niter, indices, init_variables, selected_figs=['loss', 'forward', 'probe_r_amp', 'probe_k_amp', 'probe_k_phase', 'pos'], collate_str='', show_fig=True, save_fig=False)[source]#

Wrapper function for most visualization function

ptyrad.plotting.model.plot_forward_pass(model, indices, dp_power, plot_raw=False, show_fig=True, pass_fig=False)[source]#

Plot the forward pass for the input torch model