def compute_time(x: layers.Layer, n: int):
print("computing time...")
if n == 1:
x_test = np.ones((2000, x.input_shape[1]))
elif n == 3:
x_test = np.ones(
(2000, x.input_shape[1], x.input_shape[2], x.input_shape[3]))
arg = tf.convert_to_tensor(x_test, dtype=tf.float32)
steps = 10000
tt = np.zeros(steps)
start_time = time.time()
for i in range(steps):
x.call(arg)
end_time = time.time()
tt[i] = end_time - start_time
start_time = end_time
return np.mean(tt), np.std(tt)
def plot_custom_layer(layer: Layer, input_shape: tuple, plot_filepath: str):
tmp_input = Input(shape=input_shape)
layer.build(input_shape)
tmp_output = layer.call(tmp_input)
tmp_model = Model(inputs=[tmp_input], outputs=[tmp_output])
return plot_model(tmp_model, to_file=plot_filepath, show_shapes=True)