moved framework code to dwc_server, set up local training file

dwc_client.py

@@ -0,0 +1,66 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+
+import tensorflow as tf
+import tensorflow_probability as tfp
+import os
+import numpy as np
+import asyncio
+from fastapi import FastAPI, WebSocket, File, UploadFile
+from fastapi.responses import HTMLResponse, FileResponse
+import websockets
+import uuid
+import glob
+
+client_id = 'A'
+model = None
+inputs = None
+outputs = None
+
+# set up lightweight bayesian neural network
+def sample_net(input_shape,
+    activation = tf.nn.relu,
+    batch_size = None):
+    def one_layer(x, dilation_rate=(1, 1, 1)):
+        x = tfpl.Convolution3DReparameterization(
+                filters,
+                kernel_size=3,
+                padding="same",
+                dilation_rate=dilation_rate,
+                activation=activation,
+                name="layer/vwnconv3d",
+            )(x)
+        x = tfkl.Activation(activation, name="layer/activation")(x)
+        return x
+
+    inputs = tf.keras.layers.input(shape = input_shape, batch_size=batch_size, name="inputs")
+    x = one_layer(inputs)
+
+    return tf.keras.Model(inputs=inputs, outputs=x)
+
+def train(inputs, outputs):
+    if model == None:
+        model = sample_net(np.shape(inputs))
+    # input shape: 4 x 1
+    _op = 'adam'
+    _loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
+    _metrics = ['accuracy']
+    model.compile(optimizer=_op, loss=_loss, metrics=_metrics)
+    model.fit(inputs=inputs, outputs=outputs, epochs=1, verbose=2)
+
+
+def most_recent_consolidated():
+    list_of_files = glob.glob('/server/*')
+    latest_file = max(list_of_files, key=os.path.getctime)
+    return latest_file
+
+# ping server side
+@app.get('/')
+async def send_weights():
+    return FileResponse('model-'+client_id+'.h5')
+
+@app.post("/")
+async def load_consolidated():
+    model.load_weights(most_recent_consolidated())
+    train()
+    return {'consolidated weights': model}

dwc_server.py

@@ -0,0 +1,88 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+
+import tensorflow as tf
+import os
+import asyncio
+from fastapi import FastAPI, WebSocket, File, UploadFile
+from fastapi.responses import HTMLResponse, FileResponse
+import websockets
+import copy
+import numpy as np
+import uuid
+
+app = FastAPI()
+
+datasets = []
+num_clients = 1
+model = None
+
+@app.get('/')
+async def root:
+    return 'root'
+
+@app.get("/{model}/")
+async def update_model(model: tf.keras.Model):
+    return FileResponse('consolidated.h5')
+
+@app.post("/weights")
+async def load_weights(image: UploadFile = File(...)): #do we still need to load a file here
+    model.load_weights(image.filename)
+    datasets.append(copy.deepcopy(model))
+    dwc_frame()
+    return {'consolidated weights': model}
+
+def load_data():
+    for root, dirs, files in os.walk(path):
+        for filename in files:
+            if filename[:6] != 'model-':
+                continue
+            model.load_weights(filename)
+            datasets.append(copy.deepcopy(model))
+
+
+async def dwc_frame():
+    #load weights
+    load_data()
+
+    # need to figure out how to only run this code after all (or most) client weights
+
+    # dwc implementation
+    # what are the priors
+    model = distributed_weight_consolidation(datasets, priors)
+
+    model.save_weights(
+        'server/consolidated-'+uuid.uuid4().__str__()+'.h5',
+        save_format = 'h5')
+    update_model(model)
+
+
+# LOOK AT THIS LATERRRRR
+
+async def distributed_weight_consolidation(models_weights, model_priors):
+    # models is a list of weights of client-models; models = [model1, model2, model3...]
+    num_layers =  int(len(model_weights[0])/2.0)
+    num_datasets  = np.shape(model_weights)[0]
+    consolidated_model = model_weights[0]
+    mean_idx = [i for i in range(0,len(model_weights[0])) if i % 2 == 0]
+    std_idx = [i for i in range(0,len(model_weights[0])) if i % 2 != 0]
+    ep = 1e-5
+    for i in range(num_layers):
+        num_1 = 0; num_2 = 0
+        den_1 = 0; den_2 = 0
+        for m in range(num_datasets):
+            model = model_weights[m]
+            prior = model_priors[m]
+            mu_s = model[mean_idx[i]]
+            mu_o = prior[mean_idx[i]]
+            sig_s = model[std_idx[i]]
+            sig_o = prior[std_idx[i]]
+            d1 = np.power(sig_s,2) + ep; d2= np.power(sig_o,2) + ep
+            num_1 = num_1 + (mu_s/d1)
+            num_2 = num_2 + (mu_o/d2)
+            den_1 = den_1 + (1.0/d1)
+            den_2 = den_2 + (1.0/d2)
+        consolidated_model[mean_idx[i]] =  (num_1 - num_2)/(den_1 -den_2)
+        consolidated_model[std_idx[i]] =  1/(den_1 -den_2)
+    return consolidated_model
+