How to use GPU for Keras learning in this example
I have the following code, for vertical federated learning where i use keras to create a simple 2 layer NN
``` batch_size = 32 learning_rate = 1e-3 epochs = 50
optimizer = keras.optimizers.Adam(learning_rate = learning_rate) loss_fn = keras.losses.SparseCategoricalCrossentropy(from_logits = False) train_acc_metric = tf.keras.metrics.SparseCategoricalAccuracy()
class Client(): def init(self, train, test, labelled): self.__trainX = train.copy() self.__testX = test.copy() self.labelled = labelled
if (labelled):
self.__trainY = self.__trainX.pop('credit_risk')
self.__testY = self.__testX.pop('credit_risk')
normalizer = preprocessing.Normalization()
normalizer.adapt(np.array(self.__trainX.loc[common_train_id]))
self.model = keras.Sequential([
normalizer,
layers.Dense(128, activation='elu', kernel_regularizer=regularizers.l2(0.01)),
layers.Dropout(0.5),
layers.Dense(128, activation='elu', kernel_regularizer=regularizers.l2(0.01)),
layers.Dropout(0.5),
layers.Dense(2),
layers.Softmax()])
def next_batch(self, index):
self.batchX = self.__trainX.loc[index]
if not self.labelled:
grads = []
self.model_output = np.zeros((len(index), 2))
for i in range(len(index)):
with tf.GradientTape() as gt:
gt.watch(self.model.trainable_weights)
output_by_example = self.model(self.batchX.iloc[i:i+1], training = True)
output_for_grad = output_by_example[:,1]
self.model_output[i] = output_by_example
grads.append(gt.gradient(output_for_grad, self.model.trainable_weights))
return grads
else:
self.batchY = self.__trainY.loc[index]
with tf.GradientTape() as self.gt:
self.gt.watch(self.model.trainable_weights)
self.model_output = self.model(self.batchX, training = True)
def cal_model(self):
return self.model_output
def predict(self, test_index):
return self.model.predict(self.__testX.loc[test_index])# + 1e-8
def test_answers(self, test_index):
if self.labelled:
return self.__testY.loc[test_index]
def batch_answers(self):
if self.labelled:
return self.batchY
def loss_and_update(self, a):
if not self.labelled:
raise AssertionError("This method can only be called by client 2")
self.prob = (a + self.model_output)/2
self.c = self.coefficient_and_update()/len(self.batchX)
return self.prob, loss_fn(self.batchY, self.prob)
def coefficient_and_update(self):
if not self.labelled:
raise AssertionError("This method can only be called by client 2")
p = self.prob[:,1]
c = (p-self.batchY)/((p)*(1-p))
with self.gt:
output = sum(c * self.model_output[:,1])/len(c)
grads = self.gt.gradient(output, self.model.trainable_weights)
optimizer.apply_gradients(zip(grads, self.model.trainable_weights))
return c
def update_with(self, grads):
optimizer.apply_gradients(zip(grads, self.model.trainable_weights))
def assemble_grad(self, partial_grads):
if not self. 338