NeuralTrack is a lightweight logging tool for deep learning training, designed to track loss and gradient updates efficiently without slowing down the training process. It provides easy-to-use logging functionality and visualization tools to help monitor model performance.
✅ Minimal Overhead – Asynchronous logging ensures training speed is not compromised.
✅ Loss Tracking – Logs individual loss components per epoch.
✅ Gradient Tracking – Captures gradient statistics (mean, median, max, min) for each layer.
✅ Visualizations – Generate loss and gradient plots with simple CLI commands.
Install NeuralTrack via pip:
pip install neuraltrack1️⃣ Logging Loss
In your training loop, use the LossLogger to track loss values:
from neuraltrack.logging.loss_logger import LossLogger
logger = LossLogger("loss_log.json")
for epoch in range(num_epochs):
logger.start_epoch()
for batch in dataloader:
output = Model(input)
loss1 = loss1(output)
loss2 = loss2(input, output)
total_loss = loss1+loss2
optimizer.zero_grad()
total_loss.backward()
optimizer.setp()
loss_dict = {"total_loss": total_loss, "loss1": loss1, "loss2":loss2}
logger.add_batch_loss(loss_dict)
logger.log_epoch_loss(epoch)2️⃣ Logging Gradients
Track gradient statistics every few epochs:
from neuraltrack.logging.gradient_logger import GradientLogger
grad_logger = GradientLogger("gradient_log.json", log_interval=10)
for epoch in range(num_epochs):
for batch in dataloader:
output = Model(input)
loss1 = loss1(output)
loss2 = loss2(input, output)
total_loss = loss1+loss2
optimizer.zero_grad()
total_loss.backward()
grad_logger.log_gradients(epoch, model)
optimizer.step()3️⃣ Plotting Loss
You can execute the loss plotting functionality from the CLI by running the following command in the terminal:
neuraltrack-plot-loss --log_path path_to_loss_log.json --show_plot --save_dir output_directoryArguments for LossPlotter:
--log_path: Path to the loss log JSON file (required).--show_plot: Optional flag to display the plot interactively (if this is passed, the plot will open in a window after saving it).--save_dir: Directory where the plot will be saved (optional; defaults to LossPlots).
Example Usage:
neuraltrack-plot-loss --log_path loss_log.json --show_plotThis command will generate a loss plot from the data in loss_log.json and display it interactively.
4️⃣ Plotting Gradients
To generate gradient plots, you can use the following command in your terminal:
neuraltrack-plot-gradient --log_path path_to_gradient_log.json --epoch 5 --show_plot --chart_type line --include_biasArguments for GradientPlotter:
- --log_path: Path to the gradient log JSON file (required).
--epoch: Specific epoch number to plot (optional; default is 1).--show_plot: Optional flag to display the plot interactively (if this is passed, the plot will open in a window after saving it).--chart_type: The type of chart to use. Can be either line or bar (optional; default is line).--include_bias: Flag to include bias layers in the plot (optional; default is False).
Example Usage:
neuraltrack-plot-gradient --log_path gradient_log.json --epoch 5 --show_plot --chart_type bar --include_biasThis will generate a bar chart for the gradient values of epoch 5 and display the plot interactively.
Both LossPlotter and GradientPlotter can also be used directly within Python scripts.
from neuraltrack.visualization.loss_plot import LossPlotter
# Initialize LossPlotter with the path to the log file and save directory
plotter = LossPlotter(log_path="loss_log.json", save_dir="LossPlots")
# Plot the losses (optionally show the plot)
plotter.plot_losses(show_plot=True)log_path: The path to your loss log JSON file.save_dir: The directory where the plot should be saved (optional; defaults to LossPlots).show_plot: If set to True, the plot will be shown interactively after it is generated.
from neuraltrack.visualization.gradient_plot import GradientPlotter
# Initialize GradientPlotter with your preferred arguments
plotter = GradientPlotter(
log_path="gradient_log.json",
save_dir="GradientPlots",
chart_type="line",
include_bias=False
)
# Plot gradients for epoch 5
plotter.plot_gradients(show_plot=True, epoch=5)log_path: Path to your gradient log JSON file.save_dir: Directory where the gradient plot will be saved (optional; defaults to GradientPlots).chart_type: Choose between 'bar' or 'line' for chart type (optional; defaults to line).show_plot: If set to True, it will display the plot interactively.epoch: Specify which epoch's gradient data to plot.include_bias: Set to True if you want to include bias layers in the gradient plot.
Now let's explain the parameters for LossLogger and GradientLogger: LossLogger Parameters:
log_path(str): The path to the JSON file where the loss data will be logged (default: "loss_log.json").epoch_start_time(class variable): The time when the current epoch started (used internally to track epoch duration).loss_data(dict): The dictionary where loss data is stored for each epoch.loss_components(dict): Holds the running loss components for each batch during an epoch (e.g., for multiple loss functions).total_loss(float): Tracks the total loss for the epoch (used internally).batch_count(int): The number of batches processed in the current epoch (used internally).lock(threading.Lock): Ensures thread safety when writing logs asynchronously.
log_path(str): The path to the JSON file where gradient data will be logged (default: "gradient_log.json").save_every(int): The number of epochs between saving gradient data to disk (default: 10).log_interval(int): The interval at which gradient data is logged (e.g., every 10 epochs).grad_data(dict): A dictionary where the gradient statistics for each epoch are stored.
📌 Add support for additional metrics like accuracy and learning rate tracking. 📌 Extend visualization options (e.g., smoothing, multi-run comparisons).
This project is licensed under the MIT License.
Let me know if you want any modifications! 🚀