From 994f1202b7cf531e0017cdd3ce5bd683e28283da Mon Sep 17 00:00:00 2001
From: "Dr. Kashif Rasul" <kashif.rasul@gmail.com>
Date: Thu, 16 Jan 2020 14:34:50 +0100
Subject: [PATCH] added initial SummaryWriter

---
 pts/trainer.py | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/pts/trainer.py b/pts/trainer.py
index c78b434..394df73 100644
--- a/pts/trainer.py
+++ b/pts/trainer.py
@@ -3,6 +3,7 @@ from typing import Any, List, NamedTuple, Optional, Union
 
 import torch
 import torch.nn as nn
+from torch.utils.tensorboard import SummaryWriter
 from tqdm import tqdm
 
 from torch.utils.data import DataLoader
@@ -35,6 +36,11 @@ class Trainer:
             net.parameters(), lr=self.learning_rate, weight_decay=self.weight_decay
         )
 
+        writer = SummaryWriter()
+
+        writer.add_graph(net)
+
+
         for epoch_no in range(self.epochs):
             # mark epoch start time
             tic = time.time()
@@ -59,6 +65,10 @@ class Trainer:
                         },
                         refresh=False,
                     )
+                    n_iter = epoch_no*self.num_batches_per_epoch + batch_no
+                    writer.add_scalar('Loss/train', loss.item(), n_iter)
+                    for name, param in net.named_parameters():
+                        writer.add_histogram(name, param.clone().cpu().data.numpy(), n_iter)
 
                     loss.backward()
                     optimizer.step()