license import { CustomCallback , Logs , nextFrame , util } from '@tensorflow/tfjs' ; import * as path from 'path' ; import * as ProgressBar from 'progress' ; import { summaryFileWriter , SummaryFileWriter } from './tensorboard' ; export const progressBarHelper = { ProgressBar , log : console . log } ; /**
 * Terminal-based progress bar callback for tf.Model.fit().
 */ export class ProgbarLogger extends CustomCallback { private numTrainBatchesPerEpoch : number ; private progressBar : ProgressBar ; private currentEpochBegin : number ; private epochDurationMillis : number ; private usPerStep : number ; private batchesInLatestEpoch : number ; private terminalWidth : number ; private readonly RENDER_THROTTLE_MS = 50 ; /**
     * Construtor of LoggingCallback.
     */ constructor ( ) { super ( { onTrainBegin : async ( logs ? ) => { const samples = this . params . samples as number ; const batchSize = this . params . batchSize as number ; const steps = this . params . steps as number ; if ( samples != null || steps != null ) { this . numTrainBatchesPerEpoch = samples != null ? Math . ceil ( samples / batchSize ) : steps ; } else { this . numTrainBatchesPerEpoch = 0 ; } } , onEpochBegin : async ( epoch , logs ? ) => { progressBarHelper . log ( `Epoch ${ epoch + 1 } / ${ this . params . epochs }` ) ; this . currentEpochBegin = util . now ( ) ; this . epochDurationMillis = null ; this . usPerStep = null ; this . batchesInLatestEpoch = 0 ; this . terminalWidth = process . stderr . columns ; } , onBatchEnd : async ( batch , logs ? ) => { this . batchesInLatestEpoch ++ ; if ( batch === 0 ) { this . progressBar = new progressBarHelper . ProgressBar ( 'eta=:eta :bar :placeholderForLossesAndMetrics' , { width : Math . floor ( 0.5 * this . terminalWidth ) , total : this . numTrainBatchesPerEpoch + 1 , head : `>` , renderThrottle : this . RENDER_THROTTLE_MS } ) ; } const maxMetricsStringLength = Math . floor ( this . terminalWidth * 0.5 - 12 ) ; const tickTokens = { placeholderForLossesAndMetrics : this . formatLogsAsMetricsContent ( logs , maxMetricsStringLength ) } ; if ( this . numTrainBatchesPerEpoch === 0 ) { this . progressBar . tick ( 0 , tickTokens ) ; } else { this . progressBar . tick ( tickTokens ) ; } await nextFrame ( ) ; if ( batch === this . numTrainBatchesPerEpoch - 1 ) { this . epochDurationMillis = util . now ( ) - this . currentEpochBegin ; this . usPerStep = this . params . samples != null ? this . epochDurationMillis / ( this . params . samples as number ) * 1e3 : this . epochDurationMillis / this . batchesInLatestEpoch * 1e3 ; } } , onEpochEnd : async ( epoch , logs ? ) => { if ( this . epochDurationMillis == null ) { this . epochDurationMillis = util . now ( ) - this . currentEpochBegin ; this . usPerStep = this . epochDurationMillis / this . batchesInLatestEpoch * 1e3 ; } this . progressBar . tick ( { placeholderForLossesAndMetrics : '' } ) ; const lossesAndMetricsString = this . formatLogsAsMetricsContent ( logs ) ; progressBarHelper . log ( `${ this . epochDurationMillis . toFixed ( 0 ) }ms ` + `${ this . usPerStep . toFixed ( 0 ) }us/step - ` + `${ lossesAndMetricsString }` ) ; await nextFrame ( ) ; } , } ) ; } private formatLogsAsMetricsContent ( logs , maxMetricsLength ? ) { let metricsContent = '' ; const keys = Object . keys ( logs ) . sort ( ) ; for ( const key of keys ) { if ( this . isFieldRelevant ( key ) ) { const value = logs [ key ] ; metricsContent += `${ key }=${ getSuccinctNumberDisplay ( value ) } ` ; } } if ( maxMetricsLength != null && metricsContent . length > maxMetricsLength ) { metricsContent = metricsContent . slice ( 0 , maxMetricsLength - 3 ) + '...' ; } return metricsContent ; } private isFieldRelevant ( key ) { return key !== 'batch' && key !== 'size' ; } } const BASE_NUM_DIGITS = 2 ; const MAX_NUM_DECIMAL_PLACES = 4 ; param x return export function getSuccinctNumberDisplay ( x ) { const decimalPlaces = getDisplayDecimalPlaces ( x ) ; return decimalPlaces > MAX_NUM_DECIMAL_PLACES ? x . toExponential ( BASE_NUM_DIGITS ) : x . toFixed ( decimalPlaces ) ; } param x return export function getDisplayDecimalPlaces ( x ) { if ( ! Number . isFinite ( x ) || x === 0 || x > 1 || x < - 1 ) { return BASE_NUM_DIGITS ; } else { return BASE_NUM_DIGITS - Math . floor ( Math . log10 ( Math . abs ( x ) ) ) ; } } export interface TensorBoardCallbackArgs { /**
     * The frequency at which loss and metric values are written to logs.
     *
     * Currently supported options are:
     *
     * - 'batch': Write logs at the end of every batch of training, in addition
     *   to the end of every epoch of training.
     * - 'epoch': Write logs at the end of every epoch of training.
     *
     * Note that writing logs too often slows down the training.
     *
     * Default: 'epoch'.
     */ updateFreq ? : 'batch' | 'epoch' ; } /**
 * Callback for logging to TensorBoard durnig training.
 *
 * Users are expected to access this class through the `tensorBoardCallback()`
 * factory method instead.
 */ export class TensorBoardCallback extends CustomCallback { private trainWriter : SummaryFileWriter ; private valWriter : SummaryFileWriter ; private batchesSeen : number ; private epochsSeen : number ; private readonly args : TensorBoardCallbackArgs ; constructor ( readonly logdir = './logs' , args ? ) { super ( { onBatchEnd : async ( batch , logs ? ) => { this . batchesSeen ++ ; if ( this . args . updateFreq !== 'epoch' ) { this . logMetrics ( logs , 'batch_' , this . batchesSeen ) ; } } , onEpochEnd : async ( epoch , logs ? ) => { this . epochsSeen ++ ; this . logMetrics ( logs , 'epoch_' , this . epochsSeen ) ; } , onTrainEnd : async ( logs ? ) => { if ( this . trainWriter != null ) { this . trainWriter . flush ( ) ; } if ( this . valWriter != null ) { this . valWriter . flush ( ) ; } } } ) ; this . args = args == null ? { } : args ; if ( this . args . updateFreq == null ) { this . args . updateFreq = 'epoch' ; } util . assert ( [ 'batch' , 'epoch' ] . indexOf ( this . args . updateFreq ) !== - 1 , ( ) => `Expected updateFreq to be 'batch' or 'epoch', but got ` + `${ this . args . updateFreq }` ) ; this . batchesSeen = 0 ; this . epochsSeen = 0 ; } private logMetrics ( logs , prefix , step ) { for ( const key in logs ) { if ( key === 'batch' || key === 'size' || key === 'num_steps' ) { continue ; } const VAL_PREFIX = 'val_' ; if ( key . startsWith ( VAL_PREFIX ) ) { this . ensureValWriterCreated ( ) ; const scalarName = prefix + key . slice ( VAL_PREFIX . length ) ; this . valWriter . scalar ( scalarName , logs [ key ] , step ) ; } else { this . ensureTrainWriterCreated ( ) ; this . trainWriter . scalar ( `${ prefix }${ key }` , logs [ key ] , step ) ; } } } private ensureTrainWriterCreated ( ) { this . trainWriter = summaryFileWriter ( path . join ( this . logdir , 'train' ) ) ; } private ensureValWriterCreated ( ) { this . valWriter = summaryFileWriter ( path . join ( this . logdir , 'val' ) ) ; } } param logdir param args returns doc export function tensorBoard ( logdir = './logs' , args ? ) { return new TensorBoardCallback ( logdir , args ) ; }