federer
Version:
Experiments in asynchronous federated learning and decentralized learning
101 lines • 3.4 kB
JavaScript
"use strict";
Object.defineProperty(exports, "__esModule", { value: true });
exports.createModel = void 0;
const tslib_1 = require("tslib");
const tf = tslib_1.__importStar(require("@tensorflow/tfjs-node"));
const coordinator_1 = require("../../../coordinator");
/**
* Gets a TensorFlow model to use on the MNIST dataset.
*
* @param model Name of the model to get
*
* @param numberOutputClasses Number of label classes to use for the output.
* Setting this parameter is important: if we are running a small test
* experiment, where we only have digits 0 and 1, the model should output a
* one-hot vector of length two; in other words, the model should only be able
* to predict digits that are a part of the experiment. For full experiments,
* where we use the whole MNIST dataset, `numberOutputClasses` should be 10.
*/
function createModel(options) {
const model = getModel(options);
model.compile({
optimizer: coordinator_1.getOptimizer(options.optimizer),
loss: "categoricalCrossentropy",
metrics: ["accuracy"],
});
return model;
}
exports.createModel = createModel;
function getModel(options) {
switch (options.name) {
case "2NN":
return get2NNModel(options.numberOutputClasses);
case "CNN":
return getCNNModel(options.numberOutputClasses);
}
}
const IMAGE_HEIGHT = 28;
const IMAGE_WIDTH = 28;
const IMAGE_CHANNELS = 1;
/**
* The "MNIST 2NN" model from the FedAvg paper.
*
* The model is described as follows:
*
* > A simple multilayer-perceptron with 2 hidden layers with 200 units each
* > using ReLu activations (199,210 total parameters)
*/
const get2NNModel = (numOutputClasses) => tf.sequential({
layers: [
tf.layers.dense({
name: "HiddenLayer1",
// Note that MLP inputs are flat vectors
inputShape: [IMAGE_CHANNELS * IMAGE_HEIGHT * IMAGE_WIDTH],
units: 200,
activation: "relu",
}),
tf.layers.dense({
name: "HiddenLayer2",
units: 200,
activation: "relu",
}),
tf.layers.dense({
name: "OutputLayer",
units: numOutputClasses,
activation: "softmax",
}),
],
});
/**
* The "MNIST CNN" model from the FedAvg paper.
*
* The model is described as follows in the paper:
*
* > A CNN with two 5x5 convolution layers (the first with32 channels, the
* > second with 64, each followed with 2x2max pooling), a fully connected layer
* > with 512 units andReLu activation, and a final softmax output layer
* > (1,663,370 total parameters).
*/
const getCNNModel = (numOutputClasses) => tf.sequential({
layers: [
tf.layers.conv2d({
inputShape: [IMAGE_WIDTH, IMAGE_HEIGHT, IMAGE_CHANNELS],
filters: 32,
kernelSize: 5,
activation: "relu",
padding: "same",
}),
tf.layers.maxPooling2d({ poolSize: [2, 2] }),
tf.layers.conv2d({
filters: 64,
kernelSize: 5,
activation: "relu",
padding: "same",
}),
tf.layers.maxPooling2d({ poolSize: [2, 2] }),
tf.layers.flatten(),
tf.layers.dense({ units: 512, activation: "relu" }),
tf.layers.dense({ units: numOutputClasses, activation: "softmax" }),
],
});
//# sourceMappingURL=model.js.map