# import model base and layers
from keras.models import Sequential
from keras.layers import Dense

def twoLayerFeedForward():
    # stack the layers
    clf = Sequential()
    clf.add(Dense(9, activation='relu', input_dim=3))
    clf.add(Dense(9, activation='relu'))
    clf.add(Dense(3, activation='softmax'))
    # compile the model
    clf.compile(
        loss='categorical_crossentropy', optimizer=SGD(),
        metrics=["accuracy"]
    )
    return clf

# initialize the model object
model = twoLayerFeedForward()

# call fit to train the model
# notice how hyper-parameters are set at fit, not at init
model.fit(
    X, y, epochs=50, batch_size=256,
    validation_data=(X_test, X_test)
)

# call predict to get predictions
y_pred = model.predict(X)