Logs training loss in Jupyter/Colab notebooks

jslogger.py

"""
Author: ntt123
Date: Sep 18, 2020

Usage:

!git clone https://gist.github.com/NTT123/4596e5533e573c8ceab2f319ab5d36a2 jslog
from jslog.jslogger import JSLogger
import time

logger = JSLogger('train/test loss', ['train', 'test'])

# run the following lines in the same cell
logger.show()
for i in range(1000):
    logger.log(i, {'train': i+1, 'test': i-2})
    time.sleep(5)
"""

import IPython.display
from IPython.display import display
import uuid


class JSLogger:
    def __init__(self, title, lines):
        self.title = title
        self.lines = lines
        self.dic = dict(zip(lines, range(len(lines))))
        self.value = [0.] * len(lines)
        self.colors = ['#4dc9f6', '#f67019', '#f53794', '#537bc4',
                       '#acc236', '#166a8f', '#00a950', '#58595b', '#8549ba']
        self.myid = str(uuid.uuid4().fields[0])

    def show(self):
        d = ["""
            {{
                data: [], 
                label: "{}",
                borderColor: "{}",
                borderWidth: 1.5, 
                pointRadius: 0, 
                lineTension: 0,
                backgroundColor: 'rgba(0, 0, 0, 0)',
                fill: false
            }}
            """.format(s, c) for s, c in zip(self.lines, self.colors)]
        data = ', '.join(d)

        display(IPython.display.HTML(f'''
            <script src="https://cdn.jsdelivr.net/npm/chart.js@2.9.3/dist/Chart.min.js"></script>
            <div class="chart-container" style="height:250px; width:500px;"><canvas id="{self.myid}"></canvas></div>
            <script>
            var ctx_{self.myid} = document.getElementById('{self.myid}').getContext('2d');
            ctx_{self.myid}.canvas.width = 500;
		    ctx_{self.myid}.canvas.height = 250;

            var myLineChart_{self.myid} = new Chart(ctx_{self.myid}, {{
            type: 'line',
            data: {{labels: [],datasets: [{data}]}},
            options: {{
                tooltips: {{intersect: false,mode: 'index'}},
				animation: {{duration: 0}},
                title: {{display: true,text: '{self.title}'}} }} }});
            function addData(chart, label, data) {{
                chart.data.labels.push(label);
                chart.data.datasets.forEach((dataset, i) => {{
                    dataset.data.push(data[i]);
                }});
                chart.update();
            }}
            </script>
        '''))

    def log(self, step, dic):
        for k, v in dic.items():
            i = self.dic[k]
            self.value[i] = v

        display(IPython.display.Javascript(f'''
            addData(myLineChart_{self.myid}, {step}, {self.value});    
        '''))