@epfml/discojs 2.1.2-p20240718132634.0 → 2.1.2-p20240723143623.0

package/dist/default_tasks/mnist.js

@@ -20,13 +20,14 @@ export const mnist = {
             trainingInformation: {
                 modelID: 'mnist-model',
                 epochs: 20,
-                roundDuration: 10,
+                roundDuration: 2,
                 validationSplit: 0.2,
-                batchSize: 30,
+                batchSize: 64,
                 dataType: 'image',
                 IMAGE_H: 28,
                 IMAGE_W: 28,
-                preprocessingFunctions: [data.ImagePreprocessing.Normalize],
+                // Images should already be at the right size but resizing just in case
+                preprocessingFunctions: [data.ImagePreprocessing.Resize, data.ImagePreprocessing.Normalize],
                 LABEL_LIST: ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9'],
                 scheme: 'decentralized',
                 noiseScale: undefined,
@@ -39,22 +40,24 @@ export const mnist = {
         };
     },
     getModel() {
+        // Architecture from the PyTorch MNIST example (I made it slightly smaller, 650kB instead of 5MB)
+        // https://github.com/pytorch/examples/blob/main/mnist/main.py
         const model = tf.sequential();
         model.add(tf.layers.conv2d({
             inputShape: [28, 28, 3],
-            kernelSize: 3,
-            filters: 16,
-            activation: 'relu'
+            kernelSize: 5,
+            filters: 8,
+            activation: 'relu',
         }));
+        model.add(tf.layers.conv2d({ kernelSize: 5, filters: 16, activation: 'relu' }));
         model.add(tf.layers.maxPooling2d({ poolSize: 2, strides: 2 }));
-        model.add(tf.layers.conv2d({ kernelSize: 3, filters: 32, activation: 'relu' }));
-        model.add(tf.layers.maxPooling2d({ poolSize: 2, strides: 2 }));
-        model.add(tf.layers.conv2d({ kernelSize: 3, filters: 32, activation: 'relu' }));
-        model.add(tf.layers.flatten({}));
-        model.add(tf.layers.dense({ units: 64, activation: 'relu' }));
+        model.add(tf.layers.dropout({ rate: 0.25 }));
+        model.add(tf.layers.flatten());
+        model.add(tf.layers.dense({ units: 32, activation: 'relu' }));
+        model.add(tf.layers.dropout({ rate: 0.25 }));
         model.add(tf.layers.dense({ units: 10, activation: 'softmax' }));
         model.compile({
-            optimizer: 'rmsprop',
+            optimizer: 'adam',
             loss: 'categoricalCrossentropy',
             metrics: ['accuracy']
         });

package/dist/task/task_handler.js

@@ -13,8 +13,16 @@ export async function pushTask(url, task, model) {
 export async function fetchTasks(url) {
     const response = await axios.get(new URL(TASK_ENDPOINT, url).href);
     const tasks = response.data;
-    if (!(Array.isArray(tasks) && tasks.every(isTask))) {
-        throw new Error('invalid tasks response');
+    if (!Array.isArray(tasks)) {
+        throw new Error('Expected to receive an array of Tasks when fetching tasks');
+    }
+    else if (!tasks.every(isTask)) {
+        for (const task of tasks) {
+            if (!isTask(task)) {
+                console.error("task has invalid format:", task);
+            }
+        }
+        throw new Error('invalid tasks response, the task object received is not well formatted');
     }
     return Map(tasks.map((t) => [t.id, t]));
 }

package/dist/task/training_information.d.ts

@@ -1,4 +1,3 @@
-import type { AggregatorChoice } from '../aggregator/get.js';
 import type { Preprocessing } from '../dataset/data/preprocessing/index.js';
 import { PreTrainedTokenizer } from '@xenova/transformers';
 export interface TrainingInformation {
@@ -20,7 +19,7 @@ export interface TrainingInformation {
     decentralizedSecure?: boolean;
     maxShareValue?: number;
     minimumReadyPeers?: number;
-    aggregator?: AggregatorChoice;
+    aggregator?: 'mean' | 'secure';
     tokenizer?: string | PreTrainedTokenizer;
     maxSequenceLength?: number;
     tensorBackend: 'tfjs' | 'gpt';

package/dist/task/training_information.js

@@ -19,7 +19,7 @@ export function isTrainingInformation(raw) {
         typeof validationSplit !== 'number' ||
         (tokenizer !== undefined && typeof tokenizer !== 'string' && !(tokenizer instanceof PreTrainedTokenizer)) ||
         (maxSequenceLength !== undefined && typeof maxSequenceLength !== 'number') ||
-        (aggregator !== undefined && typeof aggregator !== 'number') ||
+        (aggregator !== undefined && typeof aggregator !== 'string') ||
         (clippingRadius !== undefined && typeof clippingRadius !== 'number') ||
         (decentralizedSecure !== undefined && typeof decentralizedSecure !== 'boolean') ||
         (maxShareValue !== undefined && typeof maxShareValue !== 'number') ||
@@ -33,6 +33,13 @@ export function isTrainingInformation(raw) {
         (preprocessingFunctions !== undefined && !Array.isArray(preprocessingFunctions))) {
         return false;
     }
+    if (aggregator !== undefined) {
+        switch (aggregator) {
+            case 'mean': break;
+            case 'secure': break;
+            default: return false;
+        }
+    }
     switch (dataType) {
         case 'image': break;
         case 'tabular': break;

package/dist/training/disco.d.ts

@@ -1,6 +1,6 @@
 import { BatchLogs, data, EpochLogs, Logger, Memory, Task, TrainingInformation } from '../index.js';
 import { client as clients } from '../index.js';
-import type { Aggregator } from '../aggregator/index.js';
+import { type Aggregator } from '../aggregator/index.js';
 import type { RoundLogs } from './trainer/trainer.js';
 export interface DiscoOptions {
     client?: clients.Client;
@@ -20,7 +20,7 @@ export declare class Disco {
     readonly logger: Logger;
     readonly memory: Memory;
     private readonly client;
-    private readonly trainer;
+    private readonly trainerPromise;
     constructor(task: Task, options: DiscoOptions);
     /** Train on dataset, yielding logs of every round. */
     trainByRound(dataTuple: data.DataSplit): AsyncGenerator<RoundLogs & {

package/dist/training/disco.js

@@ -1,7 +1,7 @@
 import { List } from 'immutable';
 import { async_iterator } from '../index.js';
 import { client as clients, EmptyMemory, ConsoleLogger } from '../index.js';
-import { MeanAggregator } from '../aggregator/mean.js';
+import { getAggregator } from '../aggregator/index.js';
 import { enumerate, split } from '../utils/async_iterator.js';
 import { TrainerBuilder } from './trainer/trainer_builder.js';
 /**
@@ -14,36 +14,23 @@ export class Disco {
     logger;
     memory;
     client;
-    trainer;
+    trainerPromise;
     constructor(task, options) {
+        // Fill undefined options with default values
         if (options.scheme === undefined) {
             options.scheme = task.trainingInformation.scheme;
         }
-        if (options.aggregator === undefined) {
-            options.aggregator = new MeanAggregator();
-        }
         if (options.client === undefined) {
             if (options.url === undefined) {
                 throw new Error('could not determine client from given parameters');
             }
+            if (options.aggregator === undefined) {
+                options.aggregator = getAggregator(task, { scheme: options.scheme });
+            }
             if (typeof options.url === 'string') {
                 options.url = new URL(options.url);
             }
-            switch (options.scheme) {
-                case 'federated':
-                    options.client = new clients.federated.FederatedClient(options.url, task, options.aggregator);
-                    break;
-                case 'decentralized':
-                    options.client = new clients.decentralized.DecentralizedClient(options.url, task, options.aggregator);
-                    break;
-                case 'local':
-                    options.client = new clients.Local(options.url, task, options.aggregator);
-                    break;
-                default: {
-                    const _ = options.scheme;
-                    throw new Error('should never happen');
-                }
-            }
+            options.client = clients.getClient(options.scheme, options.url, task, options.aggregator);
         }
         if (options.logger === undefined) {
             options.logger = new ConsoleLogger();
@@ -59,7 +46,7 @@ export class Disco {
         this.memory = options.memory;
         this.logger = options.logger;
         const trainerBuilder = new TrainerBuilder(this.memory, this.task);
-        this.trainer = trainerBuilder.build(this.client, options.scheme !== 'local');
+        this.trainerPromise = trainerBuilder.build(this.client, options.scheme !== 'local');
     }
     /** Train on dataset, yielding logs of every round. */
     async *trainByRound(dataTuple) {
@@ -107,7 +94,7 @@ export class Disco {
         const trainData = dataTuple.train.preprocess().batch();
         const validationData = dataTuple.validation?.preprocess().batch() ?? trainData;
         await this.client.connect();
-        const trainer = await this.trainer;
+        const trainer = await this.trainerPromise;
         for await (const [round, epochs] of enumerate(trainer.fitModel(trainData.dataset, validationData.dataset))) {
             yield async function* () {
                 let epochsLogs = List();
@@ -131,7 +118,7 @@ export class Disco {
                 }
                 return {
                     epochs: epochsLogs,
-                    participants: this.client.nodes.size + 1, // add ourself
+                    participants: this.client.nbOfParticipants, // already includes ourselves
                 };
             }.bind(this)();
         }
@@ -141,7 +128,7 @@ export class Disco {
      * Stops the ongoing training instance without disconnecting the client.
      */
     async pause() {
-        const trainer = await this.trainer;
+        const trainer = await this.trainerPromise;
         await trainer.stopTraining();
     }
     /**

package/dist/training/trainer/trainer.js

@@ -18,7 +18,7 @@ export class Trainer {
         this.#roundDuration = task.trainingInformation.roundDuration;
         this.#epochs = task.trainingInformation.epochs;
         if (!Number.isInteger(this.#epochs / this.#roundDuration))
-            throw new Error(`round duration doesn't divide epochs`);
+            throw new Error(`round duration ${this.#roundDuration} doesn't divide number of epochs ${this.#epochs}`);
     }
     /**
      * Request stop training to be used from the Disco instance or any class that is taking care of the trainer.

package/dist/types.d.ts

@@ -2,7 +2,5 @@ import type { Map } from 'immutable';
 import type { WeightsContainer } from './index.js';
 import type { NodeID } from './client/index.js';
 export type Path = string;
-export type MetadataKey = string;
-export type MetadataValue = string;
 export type Features = number | number[] | number[][] | number[][][] | number[][][][] | number[][][][][];
 export type Contributions = Map<NodeID, WeightsContainer>;

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@epfml/discojs",
-  "version": "2.1.2-p20240718132634.0",
+  "version": "2.1.2-p20240723143623.0",
   "type": "module",
   "main": "dist/index.js",
   "types": "dist/index.d.ts",