cifar10.ts

import * as tf from "@tensorflow/tfjs";

import type { Model, TaskProvider } from "../index.js";
import { models } from "../index.js";

import baseModel from "../models/mobileNet_v1_025_224.js";

export const cifar10: TaskProvider<"image", "decentralized"> = {
  getTask() {
    return Promise.resolve({
      id: "cifar10",
      dataType: "image",
      displayInformation: {
        title: "CIFAR10",
        summary: {
          preview:
            "CIFAR-10 is a classic image classification task, and one of the most widely used datasets for machine learning research.",
          overview:
            "The dataset contains 60,000 32x32 color images in 10 different classes: airplanes, cars, birds, cats, deer, dogs, frogs, horses, ships, and trucks. The official CIFAR-10 website can be found at https://www.cs.toronto.edu/~kriz/cifar.html . You can find a link to a sample dataset at the next step.",
        },
        model:
          "The model is a pretrained MobileNetV1 model trained in Tensorflow.js. The last output layer is replaced with a fully connected layer with softmax activation and one output neuron per CIFAR10 category. The data preprocessing reshapes images into 224x224 pixels and normalizes values between 0 and 1. The neural network is optimized via Stochastic Gradient Descent and a categorical Cross Entropy loss.",
        dataFormatInformation:
          'Images should be of .png format and of size 32x32. <br> The CSV file should start with the exact header "filename,label", and each row should contain an image filename (without extension) and its label.<br><br> For example if you have images: 0.png (of a frog) and 1.png (of a car) <br> The CSV file should be: <br>filename, label <br><br> 0, frog <br> 1, car',
        dataExample:
          "https://storage.googleapis.com/deai-313515.appspot.com/example_training_data/cifar10-example.png",
        sampleDataset: {
          link: "https://storage.googleapis.com/deai-313515.appspot.com/example_training_data.tar.gz",
          instructions:
            'Opening the link should start downloading a zip file which you can unzip. To connect the data, use the CSV option below and select the file named "cifar10-labels.csv". You can now connect the images located in the "CIFAR10" folder. Note that there are only 24 images in this sample dataset which is far too few to successfully train a machine learning model.',
        },
      },
      trainingInformation: {
        epochs: 20,
        roundDuration: 10,
        validationSplit: 0.2,
        batchSize: 10,
        IMAGE_H: 224,
        IMAGE_W: 224,
        LABEL_LIST: [
          "airplane",
          "automobile",
          "bird",
          "cat",
          "deer",
          "dog",
          "frog",
          "horse",
          "ship",
          "truck",
        ],
        scheme: "decentralized",
        maxIterations: 1,
        beta: 0.9,
        aggregationStrategy: "mean",
        privacy: {
          differentialPrivacy: {
            clippingRadius: 1,
            epsilon: 50,
            delta: 1e-5,
          },
        },
        minNbOfParticipants: 3,
        maxShareValue: 100,
        tensorBackend: "tfjs",
      },
    });
  },

  async getModel(): Promise<Model<"image">> {
    const mobilenet = await tf.loadLayersModel({
      load: async () => Promise.resolve(baseModel),
    });

    const x = mobilenet.getLayer("global_average_pooling2d_1");
    const predictions = tf.layers
      .dense({ units: 10, activation: "softmax", name: "denseModified" })
      .apply(x.output) as tf.SymbolicTensor;

    const model = tf.model({
      inputs: mobilenet.input,
      outputs: predictions,
      name: "modelModified",
    });

    model.compile({
      optimizer: "sgd",
      loss: "categoricalCrossentropy",
      metrics: ["accuracy"],
    });

    return new models.TFJS("image", model);
  },
};