elems_equal = tf.equal(mask, False)
as_int = tf.cast(elems_equal, tf.int32)
count = tf.reduce_sum(as_int)
#print("COUNT K:", count)
np_u = np.array(u[i])
top_k_grad[i] = np.where(mask, 0.0, np_u)
# Send gradients to server
q = Quantize()
# Set the bitwidth here
q.bitwidth = 32
q_gradient_list = []
for each_array in top_k_grad:
q_w = q.quantize(each_array)
q_gradient_list.append(q_w)
for i in range(len(np_u)):
# Feedback error correction
r[i] = u[i] - q_gradient_list[i]
comm.send(q_gradient_list, dest=0, tag=11)
## NOT WORKING: Receive and set weights from server
#weights = comm.recv(source=0, tag=11)
#model.set_weights(weights)
## WORK AROUND: Receive and apply gradients
grad_rx = comm.recv(source=0, tag=11)
optimizer.learning_rate = 1
with tf.GradientTape() as tape:
y_pred = model(x_batch_train, training=True)
loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(
from_logits=False)
loss = loss_fn(y_batch_train, y_pred)
train_accuracy(y_batch_train, y_pred)
train_loss(loss)
#print("Step", step, loss)
gradient_list = tape.gradient(loss, model.trainable_weights)
q = Quantize()
# Set the bitwidth here
q.bitwidth = 16
q_gradient_list = []
for each_array in gradient_list:
q_w = q.quantize(each_array.numpy())
q_gradient_list.append(tf.convert_to_tensor(q_w))
# TEST
'''
for each in range(len(q_gradient_list)):
print(q_gradient_list[each])
print("+++++++++++++++++++++++++++++")
print(gradient_list[each])
sys.exit()
'''
# Send QUANTIZED gradients to server
comm.send(q_gradient_list, dest=0, tag=11)
# Send gradients to server
if step == 0:
# Compute and write grad var
std = tf.math.reduce_std(concat_grads, 0)
with grad_var_writer.as_default():
tf.summary.scalar("grad_var", std, step=epoch)
np_u = np.array(u[i])
q = Quantize()
# Set the bitwidth here
q.bitwidth = 32
q_np_u = []
for each_idx in range(len(np_u)):
q_w = q.quantize(np_u[each_idx])
q_np_u.append(q_w)
# Feedback error correction
r[each_idx] = u[each_idx] - q_np_u[each_idx]
# Send gradients to server
comm.send(r, dest=0, tag=11)
## NOT WORKING: Receive and set weights from server
#weights = comm.recv(source=0, tag=11)
#model.set_weights(weights)
## WORK AROUND: Receive and apply gradients
grad_rx = comm.recv(source=0, tag=11)