diff --git a/dwc_client.py b/dwc_client.py new file mode 100755 index 0000000..3ec6cbd --- /dev/null +++ b/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} diff --git a/dwc_server.py b/dwc_server.py new file mode 100755 index 0000000..a920088 --- /dev/null +++ b/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 + \ No newline at end of file