@epfml/discojs 0.0.1

package/dist/tasks/titanic.d.ts

@@ -0,0 +1,4 @@
+import * as tf from '@tensorflow/tfjs';
+import { Task } from '..';
+export declare const task: Task;
+export declare function model(): tf.LayersModel;

package/dist/tasks/titanic.js

@@ -0,0 +1,86 @@
+"use strict";
+Object.defineProperty(exports, "__esModule", { value: true });
+exports.model = exports.task = void 0;
+var tslib_1 = require("tslib");
+var tf = (0, tslib_1.__importStar)(require("@tensorflow/tfjs"));
+exports.task = {
+    taskID: 'titanic',
+    displayInformation: {
+        taskTitle: 'Titanic',
+        summary: "Test our platform by using a publicly available <b>tabular</b> dataset. <br><br> Download the passenger list from the Titanic shipwreck <a class='underline text-primary-dark dark:text-primary-light' href='https://www.kaggle.com/c/titanic'>here</a>. <br> This model predicts the type of person most likely to survive/die based on their characteristics (sex, age, class etc.).",
+        overview: 'We all know the unfortunate story of the Titanic: this flamboyant new transatlantic boat that sunk in 1912 in the North Atlantic Ocean. Today, we revist this tragedy by trying to predict the survival odds of the passenger given some basic features.',
+        model: 'The current form of the model does not normalize the given data and applies a very simple pre-processing of the data.',
+        tradeoffs: 'We are using a small model for this task: 2 connected layers with few nodes and we are using no optimization techniques implemented. This allows fast training but can yield to poor performance.',
+        dataFormatInformation: 'This model takes as input a CSV file with 12 columns.  general information about the passenger (sex, age, name ...) and specific related Titanic data such as the ticket class bought by the passenger, its cabin number ... \n pclass: A proxy for socio-economic status (SES) \n 1st = Upper \n 2nd = Middle \n 3rd = Lower \n \n age: Age is fractional if less than 1. If the age is estimated, is it in the form of xx.5 \n \n sibsp: The dataset defines family relations in this way... \n Sibling = brother, sister, stepbrother, stepsister \n Spouse = husband, wife (mistresses and fiancés were ignored) \n \n parch: The dataset defines family relations in this way.. \n Parent = mother, father \n Child = daughter, son, stepdaughter, stepson \n Some children travelled only with a nanny, therefore parch=0 for them. <br><br> The first line of the csv contains the header: <br> PassengerId, Survived, Pclass, Name, Sex, Age, SibSp, Parch, Ticket, Fare, Cabin, Embarked <br> Each susequent row contains the corresponding data.',
+        dataExampleText: 'Below one can find an example of a datapoint taken as input by our model. In this datapoint, the person is young man named Owen Harris that unfortunnalty perished with the Titanic. He boarded the boat in South Hamptons and was a 3rd class passenger.On the testing page, the data should not contain the class column (Survived).',
+        dataExample: [
+            { columnName: 'PassengerId', columnData: '1' },
+            { columnName: 'Survived', columnData: '0' },
+            { columnName: 'Name', columnData: 'Braund, Mr. Owen Harris' },
+            { columnName: 'Sex', columnData: 'male' },
+            { columnName: 'Age', columnData: '22' },
+            { columnName: 'SibSp', columnData: '1' },
+            { columnName: 'Parch', columnData: '0' },
+            { columnName: 'Ticket', columnData: '1/5 21171' },
+            { columnName: 'Fare', columnData: '7.25' },
+            { columnName: 'Cabin', columnData: 'E46' },
+            { columnName: 'Embarked', columnData: 'S' },
+            { columnName: 'Pclass', columnData: '3' }
+        ],
+        headers: [
+            'PassengerId',
+            'Survived',
+            'Name',
+            'Sex',
+            'Age',
+            'SibSp',
+            'Parch',
+            'Ticket',
+            'Fare',
+            'Cabin',
+            'Embarked',
+            'Pclass'
+        ]
+    },
+    trainingInformation: {
+        modelID: 'titanic-model',
+        epochs: 20,
+        roundDuration: 10,
+        validationSplit: 0.2,
+        batchSize: 30,
+        preprocessFunctions: [],
+        modelCompileData: {
+            optimizer: 'rmsprop',
+            loss: 'binaryCrossentropy',
+            metrics: ['accuracy']
+        },
+        receivedMessagesThreshold: 1,
+        dataType: 'tabular',
+        inputColumns: [
+            'PassengerId',
+            'Age',
+            'SibSp',
+            'Parch',
+            'Fare',
+            'Pclass'
+        ],
+        outputColumns: [
+            'Survived'
+        ],
+        scheme: 'Federated'
+    }
+};
+function model() {
+    var model = tf.sequential();
+    model.add(tf.layers.dense({
+        inputShape: [6],
+        units: 124,
+        activation: 'relu',
+        kernelInitializer: 'leCunNormal'
+    }));
+    model.add(tf.layers.dense({ units: 64, activation: 'relu' }));
+    model.add(tf.layers.dense({ units: 32, activation: 'relu' }));
+    model.add(tf.layers.dense({ units: 1, activation: 'sigmoid' }));
+    return model;
+}
+exports.model = model;

package/dist/testing/tester.d.ts

@@ -0,0 +1,5 @@
+import { ModelActor } from '../model_actor';
+export declare class Tester extends ModelActor {
+    testModel(): Promise<boolean>;
+}
+export default Tester;

package/dist/testing/tester.js

@@ -0,0 +1,21 @@
+"use strict";
+Object.defineProperty(exports, "__esModule", { value: true });
+exports.Tester = void 0;
+var tslib_1 = require("tslib");
+var model_actor_1 = require("../model_actor");
+var Tester = /** @class */ (function (_super) {
+    (0, tslib_1.__extends)(Tester, _super);
+    function Tester() {
+        return _super !== null && _super.apply(this, arguments) || this;
+    }
+    Tester.prototype.testModel = function () {
+        return (0, tslib_1.__awaiter)(this, void 0, void 0, function () {
+            return (0, tslib_1.__generator)(this, function (_a) {
+                return [2 /*return*/, true];
+            });
+        });
+    };
+    return Tester;
+}(model_actor_1.ModelActor));
+exports.Tester = Tester;
+exports.default = Tester;

package/dist/training/disco.d.ts

@@ -0,0 +1,12 @@
+import { Memory } from '@/memory';
+import { Client, dataset, Logger, Task, TrainingInformant, TrainingSchemes } from '..';
+export declare class Disco {
+    readonly task: Task;
+    readonly logger: Logger;
+    readonly memory: Memory;
+    private readonly client;
+    private readonly trainer;
+    constructor(task: Task, logger: Logger, memory: Memory, scheme: TrainingSchemes, informant: TrainingInformant, client: Client);
+    startTraining(data: dataset.Data): Promise<void>;
+    stopTraining(): Promise<void>;
+}

package/dist/training/disco.js

@@ -0,0 +1,62 @@
+"use strict";
+Object.defineProperty(exports, "__esModule", { value: true });
+exports.Disco = void 0;
+var tslib_1 = require("tslib");
+var __1 = require("..");
+var trainer_builder_1 = require("./trainer/trainer_builder");
+// Handles the training loop, server communication & provides the user with feedback.
+var Disco = /** @class */ (function () {
+    // client need to be connected
+    function Disco(task, logger, memory, scheme, informant, client) {
+        this.task = task;
+        this.logger = logger;
+        this.memory = memory;
+        this.client = client;
+        if (client.task !== task) {
+            throw new Error('client not setup for given task');
+        }
+        if (informant.taskID !== task.taskID) {
+            throw new Error('informant not setup for given task');
+        }
+        var trainerBuilder = new trainer_builder_1.TrainerBuilder(this.memory, this.task, informant);
+        this.trainer = trainerBuilder.build(this.client, scheme !== __1.TrainingSchemes.LOCAL);
+    }
+    Disco.prototype.startTraining = function (data) {
+        return (0, tslib_1.__awaiter)(this, void 0, void 0, function () {
+            return (0, tslib_1.__generator)(this, function (_a) {
+                switch (_a.label) {
+                    case 0:
+                        if (data.size === 0) {
+                            this.logger.error('Training aborted. No uploaded file given as input.');
+                            throw new Error('No data in dataset');
+                        }
+                        this.logger.success('Thank you for your contribution. Data preprocessing has started');
+                        return [4 /*yield*/, this.trainer];
+                    case 1: return [4 /*yield*/, (_a.sent()).trainModel(data.dataset)];
+                    case 2:
+                        _a.sent();
+                        return [2 /*return*/];
+                }
+            });
+        });
+    };
+    // Stops the training function
+    //
+    // do not disconnect given client
+    Disco.prototype.stopTraining = function () {
+        return (0, tslib_1.__awaiter)(this, void 0, void 0, function () {
+            return (0, tslib_1.__generator)(this, function (_a) {
+                switch (_a.label) {
+                    case 0: return [4 /*yield*/, this.trainer];
+                    case 1: return [4 /*yield*/, (_a.sent()).stopTraining()];
+                    case 2:
+                        _a.sent();
+                        this.logger.success('Training was successfully interrupted.');
+                        return [2 /*return*/];
+                }
+            });
+        });
+    };
+    return Disco;
+}());
+exports.Disco = Disco;

package/dist/training/index.d.ts

@@ -0,0 +1,2 @@
+export { Disco } from './disco';
+export { TrainingSchemes } from './training_schemes';

package/dist/training/index.js

@@ -0,0 +1,7 @@
+"use strict";
+Object.defineProperty(exports, "__esModule", { value: true });
+exports.TrainingSchemes = exports.Disco = void 0;
+var disco_1 = require("./disco");
+Object.defineProperty(exports, "Disco", { enumerable: true, get: function () { return disco_1.Disco; } });
+var training_schemes_1 = require("./training_schemes");
+Object.defineProperty(exports, "TrainingSchemes", { enumerable: true, get: function () { return training_schemes_1.TrainingSchemes; } });

package/dist/training/trainer/distributed_trainer.d.ts

@@ -0,0 +1,21 @@
+import * as tf from '@tensorflow/tfjs';
+import { Client, Memory, Task, TrainingInformant } from '@/.';
+import { Trainer } from './trainer';
+/**
+ * Class whose role is to train a model in a distributed way with a given dataset.
+ */
+export declare class DistributedTrainer extends Trainer {
+    private readonly previousRoundModel;
+    private readonly client;
+    /** DistributedTrainer constructor, accepts same arguments as Trainer and in additional also a client who takes care of communicating weights.
+     */
+    constructor(task: Task, trainingInformant: TrainingInformant, memory: Memory, model: tf.LayersModel, previousRoundModel: tf.LayersModel, client: Client);
+    /**
+     * Callback called every time a round is over
+     */
+    onRoundEnd(accuracy: number): Promise<void>;
+    /**
+     * Callback called once training is over
+     */
+    onTrainEnd(): Promise<void>;
+}

package/dist/training/trainer/distributed_trainer.js

@@ -0,0 +1,60 @@
+"use strict";
+Object.defineProperty(exports, "__esModule", { value: true });
+exports.DistributedTrainer = void 0;
+var tslib_1 = require("tslib");
+var trainer_1 = require("./trainer");
+/**
+ * Class whose role is to train a model in a distributed way with a given dataset.
+ */
+var DistributedTrainer = /** @class */ (function (_super) {
+    (0, tslib_1.__extends)(DistributedTrainer, _super);
+    /** DistributedTrainer constructor, accepts same arguments as Trainer and in additional also a client who takes care of communicating weights.
+     */
+    function DistributedTrainer(task, trainingInformant, memory, model, previousRoundModel, client) {
+        var _this = _super.call(this, task, trainingInformant, memory, model) || this;
+        _this.previousRoundModel = previousRoundModel;
+        _this.client = client;
+        return _this;
+    }
+    /**
+     * Callback called every time a round is over
+     */
+    DistributedTrainer.prototype.onRoundEnd = function (accuracy) {
+        return (0, tslib_1.__awaiter)(this, void 0, void 0, function () {
+            var currentRoundWeights, previousRoundWeights, aggregatedWeights;
+            return (0, tslib_1.__generator)(this, function (_a) {
+                switch (_a.label) {
+                    case 0:
+                        currentRoundWeights = this.model.weights.map(function (w) { return w.read(); });
+                        previousRoundWeights = this.previousRoundModel.weights.map(function (w) { return w.read(); });
+                        return [4 /*yield*/, this.client.onRoundEndCommunication(currentRoundWeights, previousRoundWeights, this.roundTracker.round, this.trainingInformant)];
+                    case 1:
+                        aggregatedWeights = _a.sent();
+                        this.previousRoundModel.setWeights(currentRoundWeights);
+                        this.model.setWeights(aggregatedWeights);
+                        return [4 /*yield*/, this.memory.updateWorkingModel(this.task.taskID, this.trainingInformation.modelID, this.model)];
+                    case 2:
+                        _a.sent();
+                        return [2 /*return*/];
+                }
+            });
+        });
+    };
+    /**
+     * Callback called once training is over
+     */
+    DistributedTrainer.prototype.onTrainEnd = function () {
+        return (0, tslib_1.__awaiter)(this, void 0, void 0, function () {
+            return (0, tslib_1.__generator)(this, function (_a) {
+                switch (_a.label) {
+                    case 0: return [4 /*yield*/, this.client.onTrainEndCommunication(this.model.weights.map(function (w) { return w.read(); }), this.trainingInformant)];
+                    case 1:
+                        _a.sent();
+                        return [2 /*return*/];
+                }
+            });
+        });
+    };
+    return DistributedTrainer;
+}(trainer_1.Trainer));
+exports.DistributedTrainer = DistributedTrainer;

package/dist/training/trainer/local_trainer.d.ts

@@ -0,0 +1,10 @@
+import { Trainer } from './trainer';
+/** Class whose role is to locally train a model with a given dataset.
+ */
+export declare class LocalTrainer extends Trainer {
+    /**
+     * Callback called every time a round is over
+     */
+    onRoundEnd(accuracy: number): Promise<void>;
+    onTrainEnd(): Promise<void>;
+}

package/dist/training/trainer/local_trainer.js

@@ -0,0 +1,37 @@
+"use strict";
+Object.defineProperty(exports, "__esModule", { value: true });
+exports.LocalTrainer = void 0;
+var tslib_1 = require("tslib");
+var trainer_1 = require("./trainer");
+/** Class whose role is to locally train a model with a given dataset.
+ */
+var LocalTrainer = /** @class */ (function (_super) {
+    (0, tslib_1.__extends)(LocalTrainer, _super);
+    function LocalTrainer() {
+        return _super !== null && _super.apply(this, arguments) || this;
+    }
+    /**
+     * Callback called every time a round is over
+     */
+    LocalTrainer.prototype.onRoundEnd = function (accuracy) {
+        return (0, tslib_1.__awaiter)(this, void 0, void 0, function () {
+            return (0, tslib_1.__generator)(this, function (_a) {
+                switch (_a.label) {
+                    case 0: return [4 /*yield*/, this.memory.updateWorkingModel(this.task.taskID, this.trainingInformation.modelID, this.model)];
+                    case 1:
+                        _a.sent();
+                        return [2 /*return*/];
+                }
+            });
+        });
+    };
+    LocalTrainer.prototype.onTrainEnd = function () {
+        return (0, tslib_1.__awaiter)(this, void 0, void 0, function () {
+            return (0, tslib_1.__generator)(this, function (_a) {
+                return [2 /*return*/];
+            });
+        });
+    };
+    return LocalTrainer;
+}(trainer_1.Trainer));
+exports.LocalTrainer = LocalTrainer;

package/dist/training/trainer/round_tracker.d.ts

@@ -0,0 +1,30 @@
+/**
+ * Class that keeps track of the current batch in order for the trainer to query when round has ended.
+ *
+ * @remark
+ * In distributed training, the client trains locally for a certain amount of epochs before sharing his weights to the server/neighbor, this
+ * is what we call a round.
+ *
+ * The role of the RoundTracker is to keep track of when a roundHasEnded using the current batch number. The batch in the RoundTracker is cumulative whereas
+ * in the onBatchEnd it is not (it resets to 0 after each epoch).
+ *
+ * The roundDuration is the length of a round (in batches).
+ */
+export declare class RoundTracker {
+    round: number;
+    batch: number;
+    roundDuration: number;
+    constructor(roundDuration: number);
+    /**
+     * Update the batch number, to be called inside onBatchEnd. (We do not use batch output of onBatchEnd since it is
+     * not cumulative).
+     */
+    updateBatch(): void;
+    /**
+     * Returns true if a local round has ended, false otherwise.
+     *
+     * @remark
+     * Returns true if (batch) mod (batches per round) == 0, false otherwise
+     */
+    roundHasEnded(): boolean;
+}

package/dist/training/trainer/round_tracker.js

@@ -0,0 +1,44 @@
+"use strict";
+Object.defineProperty(exports, "__esModule", { value: true });
+exports.RoundTracker = void 0;
+/**
+ * Class that keeps track of the current batch in order for the trainer to query when round has ended.
+ *
+ * @remark
+ * In distributed training, the client trains locally for a certain amount of epochs before sharing his weights to the server/neighbor, this
+ * is what we call a round.
+ *
+ * The role of the RoundTracker is to keep track of when a roundHasEnded using the current batch number. The batch in the RoundTracker is cumulative whereas
+ * in the onBatchEnd it is not (it resets to 0 after each epoch).
+ *
+ * The roundDuration is the length of a round (in batches).
+ */
+var RoundTracker = /** @class */ (function () {
+    function RoundTracker(roundDuration) {
+        this.round = 0;
+        this.batch = 0;
+        this.roundDuration = roundDuration;
+    }
+    /**
+     * Update the batch number, to be called inside onBatchEnd. (We do not use batch output of onBatchEnd since it is
+     * not cumulative).
+     */
+    RoundTracker.prototype.updateBatch = function () {
+        this.batch += 1;
+    };
+    /**
+     * Returns true if a local round has ended, false otherwise.
+     *
+     * @remark
+     * Returns true if (batch) mod (batches per round) == 0, false otherwise
+     */
+    RoundTracker.prototype.roundHasEnded = function () {
+        if (this.batch === 0) {
+            return false;
+        }
+        var roundHasEnded = this.batch % this.roundDuration === 0;
+        return roundHasEnded;
+    };
+    return RoundTracker;
+}());
+exports.RoundTracker = RoundTracker;

package/dist/training/trainer/trainer.d.ts

@@ -0,0 +1,66 @@
+import * as tf from '@tensorflow/tfjs';
+import { Memory, Task, TrainingInformant, TrainingInformation } from '@/.';
+import { RoundTracker } from './round_tracker';
+import { TrainerLog } from '../../logging/trainer_logger';
+/** Abstract class whose role is to train a model with a given dataset. This can be either done
+ * locally or in a distributed way. The Trainer works as follows:
+ *
+ * 1. Call trainModel(dataset) to start training
+ * 2. Once a batch ends, onBatchEnd is triggered, witch will then call onRoundEnd once the round has ended.
+ *
+ * The onRoundEnd needs to be implemented to specify what actions to do when the round has ended. To know when
+ * a round has ended we use the roundTracker object.
+ */
+export declare abstract class Trainer {
+    readonly task: Task;
+    readonly trainingInformant: TrainingInformant;
+    readonly memory: Memory;
+    readonly model: tf.LayersModel;
+    readonly trainingInformation: TrainingInformation;
+    readonly roundTracker: RoundTracker;
+    private stopTrainingRequested;
+    private readonly trainerLogger;
+    /**
+     * Constructs the training manager.
+     * @param task the trained task
+     * @param trainingInformant the training informant
+     */
+    constructor(task: Task, trainingInformant: TrainingInformant, memory: Memory, model: tf.LayersModel);
+    /**
+     * Every time a round ends this function will be called
+     */
+    protected abstract onRoundEnd(accuracy: number): Promise<void>;
+    /**
+     * When the training ends this function will be call
+     */
+    protected abstract onTrainEnd(): Promise<void>;
+    /**
+     * Request stop training to be used from the Disco instance or any class that is taking care of the trainer.
+     */
+    stopTraining(): Promise<void>;
+    /**
+     * Start training the model with the given dataset
+     * @param dataset
+     */
+    trainModel(dataset: tf.data.Dataset<tf.TensorContainer>): Promise<void>;
+    /**
+     * Format accuracy
+     */
+    private roundDecimals;
+    /**
+     * We update the training graph, this needs to be done on epoch end as there is no validation accuracy onBatchEnd.
+     */
+    private onEpochEnd;
+    /** onBatchEnd callback, when a round ends, we call onRoundEnd (to be implemented for local and distributed instances)
+     */
+    private onBatchEnd;
+    /**
+     * reset stop training state
+     */
+    private resetStopTrainerState;
+    /**
+     * If stop training is requested, do so
+     */
+    private stopTrainModelIfRequested;
+    getTrainerLog(): TrainerLog;
+}

package/dist/training/trainer/trainer.js

@@ -0,0 +1,146 @@
+"use strict";
+Object.defineProperty(exports, "__esModule", { value: true });
+exports.Trainer = void 0;
+var tslib_1 = require("tslib");
+var round_tracker_1 = require("./round_tracker");
+var trainer_logger_1 = require("../../logging/trainer_logger");
+/** Abstract class whose role is to train a model with a given dataset. This can be either done
+ * locally or in a distributed way. The Trainer works as follows:
+ *
+ * 1. Call trainModel(dataset) to start training
+ * 2. Once a batch ends, onBatchEnd is triggered, witch will then call onRoundEnd once the round has ended.
+ *
+ * The onRoundEnd needs to be implemented to specify what actions to do when the round has ended. To know when
+ * a round has ended we use the roundTracker object.
+ */
+var Trainer = /** @class */ (function () {
+    /**
+     * Constructs the training manager.
+     * @param task the trained task
+     * @param trainingInformant the training informant
+     */
+    function Trainer(task, trainingInformant, memory, model) {
+        this.task = task;
+        this.trainingInformant = trainingInformant;
+        this.memory = memory;
+        this.model = model;
+        this.stopTrainingRequested = false;
+        this.trainerLogger = new trainer_logger_1.TrainerLogger();
+        var trainingInformation = task.trainingInformation;
+        if (trainingInformation === undefined) {
+            throw new Error('round duration is undefined');
+        }
+        this.trainingInformation = trainingInformation;
+        this.roundTracker = new round_tracker_1.RoundTracker(trainingInformation.roundDuration);
+    }
+    //* ***************************************************************
+    //* Functions to be used by the training manager
+    //* ***************************************************************
+    /**
+     * Request stop training to be used from the Disco instance or any class that is taking care of the trainer.
+     */
+    Trainer.prototype.stopTraining = function () {
+        return (0, tslib_1.__awaiter)(this, void 0, void 0, function () {
+            return (0, tslib_1.__generator)(this, function (_a) {
+                this.stopTrainingRequested = true;
+                return [2 /*return*/];
+            });
+        });
+    };
+    /**
+     * Start training the model with the given dataset
+     * @param dataset
+     */
+    Trainer.prototype.trainModel = function (dataset) {
+        return (0, tslib_1.__awaiter)(this, void 0, void 0, function () {
+            var _this = this;
+            return (0, tslib_1.__generator)(this, function (_a) {
+                switch (_a.label) {
+                    case 0:
+                        this.resetStopTrainerState();
+                        // Assign callbacks and start training
+                        return [4 /*yield*/, this.model.fitDataset(dataset.batch(this.trainingInformation.batchSize), {
+                                epochs: this.trainingInformation.epochs,
+                                validationData: dataset.batch(this.trainingInformation.batchSize),
+                                callbacks: {
+                                    onEpochEnd: function (epoch, logs) { return _this.onEpochEnd(epoch, logs); },
+                                    onBatchEnd: function (epoch, logs) { return (0, tslib_1.__awaiter)(_this, void 0, void 0, function () { return (0, tslib_1.__generator)(this, function (_a) {
+                                        switch (_a.label) {
+                                            case 0: return [4 /*yield*/, this.onBatchEnd(epoch, logs)];
+                                            case 1: return [2 /*return*/, _a.sent()];
+                                        }
+                                    }); }); }
+                                }
+                            })];
+                    case 1:
+                        // Assign callbacks and start training
+                        _a.sent();
+                        return [2 /*return*/];
+                }
+            });
+        });
+    };
+    /**
+     * Format accuracy
+     */
+    Trainer.prototype.roundDecimals = function (accuracy, decimalsToRound) {
+        if (decimalsToRound === void 0) { decimalsToRound = 2; }
+        return +(accuracy * 100).toFixed(decimalsToRound);
+    };
+    /**
+     * We update the training graph, this needs to be done on epoch end as there is no validation accuracy onBatchEnd.
+     */
+    Trainer.prototype.onEpochEnd = function (epoch, logs) {
+        this.trainerLogger.onEpochEnd(epoch, logs);
+        if (logs !== undefined && !isNaN(logs.acc) && !isNaN(logs.val_acc)) {
+            this.trainingInformant.updateTrainingAccuracyGraph(this.roundDecimals(logs.acc));
+            this.trainingInformant.updateValidationAccuracyGraph(this.roundDecimals(logs.val_acc));
+        }
+        else {
+            this.trainerLogger.error('onEpochEnd: NaN value');
+        }
+    };
+    /** onBatchEnd callback, when a round ends, we call onRoundEnd (to be implemented for local and distributed instances)
+     */
+    Trainer.prototype.onBatchEnd = function (_, logs) {
+        return (0, tslib_1.__awaiter)(this, void 0, void 0, function () {
+            return (0, tslib_1.__generator)(this, function (_a) {
+                switch (_a.label) {
+                    case 0:
+                        if (logs === undefined) {
+                            return [2 /*return*/];
+                        }
+                        this.roundTracker.updateBatch();
+                        this.stopTrainModelIfRequested();
+                        if (!this.roundTracker.roundHasEnded()) return [3 /*break*/, 2];
+                        return [4 /*yield*/, this.onRoundEnd(logs.acc)];
+                    case 1:
+                        _a.sent();
+                        _a.label = 2;
+                    case 2: return [2 /*return*/];
+                }
+            });
+        });
+    };
+    /**
+     * reset stop training state
+     */
+    Trainer.prototype.resetStopTrainerState = function () {
+        this.model.stopTraining = false;
+        this.stopTrainingRequested = false;
+    };
+    /**
+     * If stop training is requested, do so
+     */
+    Trainer.prototype.stopTrainModelIfRequested = function () {
+        if (this.stopTrainingRequested) {
+            this.model.stopTraining = true;
+            this.stopTrainingRequested = false;
+        }
+    };
+    Trainer.prototype.getTrainerLog = function () {
+        return this.trainerLogger.log;
+    };
+    return Trainer;
+}());
+exports.Trainer = Trainer;

package/dist/training/trainer/trainer_builder.d.ts

@@ -0,0 +1,25 @@
+import { Client, Task, TrainingInformant, Memory } from '../..';
+import { Trainer } from './trainer';
+/**
+ * A class that helps build the Trainer and auxiliary classes.
+ */
+export declare class TrainerBuilder {
+    private readonly memory;
+    private readonly task;
+    private readonly trainingInformant;
+    constructor(memory: Memory, task: Task, trainingInformant: TrainingInformant);
+    /**
+     * Builds a trainer object.
+     *
+     * @param client client to share weights with (either distributed or federated)
+     * @param distributed whether to build a distributed or local trainer
+     * @returns
+     */
+    build(client: Client, distributed?: boolean): Promise<Trainer>;
+    /**
+     * If a model exists in memory, laod it, otherwise load model from server
+     * @returns
+     */
+    private getModel;
+    private updateModelInformation;
+}