model.py

def classification_nn_model(input_features):
	initializer = tf.compat.v1.keras.initializers.random_normal(0.0, 0.01)
	model = tf.keras.Sequential(
		  [
		 	keraslayers.Dense(
		 		512, 
		 		activation = tf.nn.tanh, 
		 		input_shape = (input_features,), 
		 		kernel_initializer = initializer,
		 		bias_initializer = 'zeros'
		 		),
		 	keraslayers.Dense(
		 		128,  
		 		activation=tf.nn.tanh,
		 		kernel_initializer = initializer,
		 		bias_initializer='zeros'
		 		),
		 	keraslayers.Dense(
		 		32, 
		 		activation=tf.nn.tanh,
		 		kernel_initializer = initializer,
		 		bias_initializer='zeros'
		 		),
		 	keraslayers.Dense(
		 		100, 
		 		kernel_initializer = initializer,
		 		bias_initializer='zeros'
		 	)
		  ]
		)
	return model
# input_features should be changed according to the model
input_features = 600
cmodelA = classification_nn_model(input_features)

# Restore weights of the target classification model 
# on which the attack model will be trained.
# `cprefix` path can be decided by the user.
cprefix = 'models/model_checkpoints/classification'
class_ckpt_dir = tf.train.latest_checkpoint(cprefix)
cmodelA.load_weights(class_ckpt_dir)