update server/client/model interfaces

[asross]

This is a WIP pull request is relative to the audio branch, which will definitely break the MNIST and CIFAR experiments. I'm making this PR separate from the audio changes / larger reorganization to solicit feedback on the API changes alone.

The main changes are:

• the client API now handles fitting the model, and details like numExamples no longer need to be managed by users
• FederatedModel now has a simpler interface and can more transparently wrap tf.Model

Just as I'm making this, I can think of a few opportunities to simplify the interface even further, but I'll elaborate in comments on the diff.

---

This change is 

[asross]

audio_demo/index.js, line 81 at r1 (raw file):

fedModel.setup().then(async () => {
  const model = fedModel.model;

This is still awkward -- although it's nice to be able to access the underlying tf.Model (and maybe we'll still need to in order to compute the input shape), we shouldn't have to for things like prediction.

---

Comments from Reviewable

[asross]

audio_demo/index.js, line 92 at r1 (raw file):

    modelVersion.innerText = msg.modelId;
  });
  await clientAPI.setup(serverURL);

Instead of calling fedModel.setup and then clientAPI.setup, maybe the client API should mimic the server API and accept an uninitialized FederatedModel, but have it's setup method initialize both itself and the federated model. So basically, this code would look like this instead:

const serverURL = 'http://foo.bar';
const model = new AudioTransferLearningModel();
const api = new ClientAPI(model, serverURL);
api.setup().then(() => {
  // rest of logic
});

---

Comments from Reviewable

[asross]

src/client/comm.ts, line 131 at r1 (raw file):

   * TODO: consider having this method save copies of `xs` and `ys` when there
   * are too few examples, and only doing training/uploading after reaching a
   * configurable threshold (disposing of the copies afterwards).

Interested in feedback on the idea laid out in this comment -- might be a nice way of delaying updates until we get sufficient data without forcing the user to deal with that complexity. Not sure what the best way to transmit that configuration is, though.

---

Comments from Reviewable

[asross]

src/client/comm.ts, line 140 at r1 (raw file):

    const modelId = this.msg.modelId;
    // fit the model to the new data
    await this.model.fit(xs, ys, this.msg.fitConfig);

One thing I realized while writing this is that ModelFitConfig doesn't let you pass options to your optimizer (like the learning rate). For now it's probably fine if that stays fixed, but it would be nice if the server could dynamically set clients' learning rates through fitConfig.

---

Comments from Reviewable

[aman-tiwari]

Review status: 0 of 1 LGTMs obtained

---

audio_demo/index.js, line 81 at r1 (raw file):

Previously, asross (Andrew Ross) wrote…

This is still awkward -- although it's nice to be able to access the underlying tf.Model (and maybe we'll still need to in order to compute the input shape), we shouldn't have to for things like prediction.

You could cleanly do this as

fedModel.setup().then(async ({ model }) => {
 ...
})

Or

fedModel.setup().then(async ({ model: unwrappedModel }) => {

etc

---

src/client/comm.ts, line 131 at r1 (raw file):

Previously, asross (Andrew Ross) wrote…

Interested in feedback on the idea laid out in this comment -- might be a nice way of delaying updates until we get sufficient data without forcing the user to deal with that complexity. Not sure what the best way to transmit that configuration is, though.

Does reverting to the same weights give us better loss (e.g so that the clients don't diverge too much if the syncs are slow/don't run ahead of each other too much)? Does this also mean if a user is sufficiently multimodal data (w/ modality correlated with time) a single user gets to act as multiple clients.

---

src/client/comm.ts, line 140 at r1 (raw file):

Previously, asross (Andrew Ross) wrote…

One thing I realized while writing this is that ModelFitConfig doesn't let you pass options to your optimizer (like the learning rate). For now it's probably fine if that stays fixed, but it would be nice if the server could dynamically set clients' learning rates through fitConfig.

Perhaps if we keep needing to add things like this it might be worth creating a FederatedOptimiser class that handles it.

---

src/client/comm.ts, line 167 at r1 (raw file):

  protected setVarsFromMessage(newVars: SerializedVariable[]) {
    tf.tidy(() => {
      this.model.setVars(newVars.map(deserializeVar));

Minor (very minor) : try .map(v => deserializeVar(v)) because if we ever change deserializeVar to take more than 1 argument it'll break here as map calls the mapper function with (value, index, array) rather than just value.

---

Comments from Reviewable

[dsmilkov]

This is really nice simplification! Left a few comments. Feel free to come chat in person (much faster) if you have any questions

---

Reviewed 7 of 8 files at r1.
Review status: 0 of 1 LGTMs obtained

---

audio_demo/index.js, line 81 at r1 (raw file):

Previously, aman-tiwari (Aman Tiwari) wrote…

You could cleanly do this as

fedModel.setup().then(async ({ model }) => {
 ...
})```

Or
```js
fedModel.setup().then(async ({ model: unwrappedModel }) => {

etc

Require that the user passes a tf.Model in the c-tor of FederatedTfModel (and therefore in the c-tor of AudioTransferLearningModel) and make FederatedTfModel.model private

---

audio_demo/index.js, line 92 at r1 (raw file):

Previously, asross (Andrew Ross) wrote…

Instead of calling fedModel.setup and then clientAPI.setup, maybe the client API should mimic the server API and accept an uninitialized FederatedModel, but have it's setup method initialize both itself and the federated model. So basically, this code would look like this instead:

const serverURL = 'http://foo.bar';
const model = new AudioTransferLearningModel();
const api = new ClientAPI(model, serverURL);
api.setup().then(() => {
  // rest of logic
});

This is really nice - removes the need for 2 setup() calls. I would do api.connect() instead of api.setup(). In this case, you don't need TF

---

src/comm_test.ts, line 111 at r1 (raw file):

  it('transmits fit config on startup', () => {
    expect(clientAPI.msg.fitConfig.batchSize).toBe(batchSize);

instead of reading the msg field (which should be private), how about calling onUpdate() and make sure it gets called with the right batchSize

---

src/comm_test.ts, line 147 at r1 (raw file):

    await clientAPI.federatedUpdate(dummyX3, dummyY);

    waitUntil(() => clientAPI.msg.modelId !== modelId, () => {

Now that you have onUpdate() and msg field being private (see previous comment), no need to do waitUntil(). Test the public API instead by subscribing via calling onUpdate()

---

src/types.ts, line 25 at r1 (raw file):

export interface FederatedModel {
  setup(): Promise<void>;

let's add some short jsdocs to these methods , since this is public api.

---

src/client/comm.ts, line 55 at r1 (raw file):

 */

type UpdateCallback = (msg: DownloadMsg) => void;

how about renaming DownloadMsg to ClientUpdateMsg to be more explicit?

---

src/client/comm.ts, line 59 at r1 (raw file):

export class ClientAPI {
  public model: FederatedModel;
  public msg: DownloadMsg;

make this msg field private? since there is no reason the user will need this. They have onUpdate() method to subscribe to these messages.

---

src/client/comm.ts, line 131 at r1 (raw file):

Previously, aman-tiwari (Aman Tiwari) wrote…

Does reverting to the same weights give us better loss (e.g so that the clients don't diverge too much if the syncs are slow/don't run ahead of each other too much)? Does this also mean if a user is sufficiently multimodal data (w/ modality correlated with time) a single user gets to act as multiple clients.

Let's leave it for now (Having the TODO there is great). Finding the right balance between magic vs user-control will be easier to find once we have a few experiments.

---

src/client/comm.ts, line 140 at r1 (raw file):

Previously, aman-tiwari (Aman Tiwari) wrote…

Perhaps if we keep needing to add things like this it might be worth creating a FederatedOptimiser class that handles it.

If I understand correctly, fitConfig is a dictionary that the user fully controls right? Meaning the user passes this obj when constructing the ServerAPI, and later gets it on the other side in the ClientAPI. Then there is no reason to explicitly constraint it. You can type it as {[key: string] => any}

---

Comments from Reviewable

[asross]

Review status: 0 of 1 LGTMs obtained

---

src/client/comm.ts, line 55 at r1 (raw file):

Previously, dsmilkov (Daniel Smilkov) wrote…

how about renaming DownloadMsg to ClientUpdateMsg to be more explicit?

So... in the changes I just pushed up, I ended up doubling down(loads); i.e. changing UpdateCallback to DownloadCallback rather than changing DownloadMsg to ClientUpdateMsg. This makes it so we consistently use Download/Upload to refer to server-to-client/client-to-server communication everywhere in the codebase, which I think addresses the problem. But if we want to gsub Download and Upload for other terms, I'd be ok with that too!

---

src/client/comm.ts, line 131 at r1 (raw file):

Previously, dsmilkov (Daniel Smilkov) wrote…

Let's leave it for now (Having the TODO there is great). Finding the right balance between magic vs user-control will be easier to find once we have a few experiments.

@aman-tiwari yep -- reverting to the same weights means that one user gets to act as multiple clients, which is already true on the server side. If we didn't revert, then users would be sending us multiple updates for the same data points, and the first data point they happened to see would be overemphasized relative to the last one.

There are compelling reasons against reverting too; e.g. if server downloads are few and far between, and we want clients to improve locally in the interim. But there are potential downsides too. I summarized some of the pros and cons in #7.

---

Comments from Reviewable

[asross]

src/serialization.ts, line 68 at r2 (raw file):

  numExamples: number,
  vars: TensorJson[],
  modelVersion?: string,

I ended up renaming modelId to modelVersion just because that's how it felt most natural to refer to it on the client side, and since it is a monotonically increasing timestamp, I think it's an appropriate name.

---

Comments from Reviewable

[asross]

Review status: 0 of 1 LGTMs obtained

---

src/types.ts, line 30 at r2 (raw file):

}

export class FederatedTfModel implements FederatedModel {

Part of me wants to rename FederatedTfModel to just Federated, so we can write code like:

const model = await tf.loadModel(url);
const fedModel = Federated(model);

But that might be too cute / less clear.

---

Comments from Reviewable

[dsmilkov]

Nice!!

---

Reviewed 1 of 8 files at r1, 9 of 11 files at r2, 2 of 2 files at r3.
Review status: complete! 1 of 1 LGTMs obtained

---

src/serialization.ts, line 68 at r2 (raw file):

Previously, asross (Andrew Ross) wrote…

I ended up renaming modelId to modelVersion just because that's how it felt most natural to refer to it on the client side, and since it is a monotonically increasing timestamp, I think it's an appropriate name.

Sounds good.

---

src/types.ts, line 30 at r2 (raw file):

Previously, asross (Andrew Ross) wrote…

Part of me wants to rename FederatedTfModel to just Federated, so we can write code like:

const model = await tf.loadModel(url);
const fedModel = Federated(model);

But that might be too cute / less clear.

Not for this PR: I can imagine having a free-floating function fromTFModel instead of a public constructor, which is very js style. For example, when this is all done:

import * as client from '@tensorflow/federated-client';
const clientApi = client.fromTFModel(model);
// The intermediate FederatedTFModel object is not exposed.
...

---

src/client/comm.ts, line 55 at r1 (raw file):

Previously, asross (Andrew Ross) wrote…

So... in the changes I just pushed up, I ended up doubling down(loads); i.e. changing UpdateCallback to DownloadCallback rather than changing DownloadMsg to ClientUpdateMsg. This makes it so we consistently use Download/Upload to refer to server-to-client/client-to-server communication everywhere in the codebase, which I think addresses the problem. But if we want to gsub Download and Upload for other terms, I'd be ok with that too!

Sounds good! +1 to having consistency between the type and the subscribing method name,

---

src/client/comm.ts, line 57 at r3 (raw file):

   * Construct a client API for federated learning that will push and pull
   * `model` updates from the server.
   * @param {model<FederatedModel>} model - model to use with federated learning

no need for closure style annotation (remove {model<FederatedModel>} after @param

---

Comments from Reviewable

[asross]

Review status: complete! 1 of 1 LGTMs obtained

---

src/types.ts, line 30 at r2 (raw file):

Previously, dsmilkov (Daniel Smilkov) wrote…

Not for this PR: I can imagine having a free-floating function fromTFModel instead of a public constructor, which is very js style. For example, when this is all done:

import * as client from '@tensorflow/federated-client';
const clientApi = client.fromTFModel(model);
// The intermediate FederatedTFModel object is not exposed.
...

Yep, that would work well! In the next PR I modified the client/server APIs to take either a FederatedModel or a tf.Model, which feels very natural as well!

---

src/client/comm.ts, line 57 at r3 (raw file):

Previously, dsmilkov (Daniel Smilkov) wrote…

no need for closure style annotation (remove {model<FederatedModel>} after @param

Sounds good. Will fix in next PR to avoid merge conflicts (which will be arriving just after this gets merged)

---

Comments from Reviewable