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)
  }
}