TypeError: softmax() got an unexpected keyword argument 'axis'

今天我裝了tensorflow-gpu 1.4, keras 2.1.6發現運行keras模型的時候出現了以下的錯誤：、

Traceback (most recent call last):
  File "models.py", line 24, in <module>
    model = NIH()
  File "models.py", line 19, in NIH
    model.add(Activation('softmax'))
  File "/home/eric/anaconda3/lib/python3.6/site-packages/keras/models.py", line 522, in add
    output_tensor = layer(self.outputs[0])
  File "/home/eric/anaconda3/lib/python3.6/site-packages/keras/engine/topology.py", line 619, in __call__
    output = self.call(inputs, **kwargs)
  File "/home/eric/anaconda3/lib/python3.6/site-packages/keras/layers/core.py", line 304, in call
    return self.activation(inputs)
  File "/home/eric/anaconda3/lib/python3.6/site-packages/keras/activations.py", line 29, in softmax
    return K.softmax(x)
  File "/home/eric/anaconda3/lib/python3.6/site-packages/keras/backend/tensorflow_backend.py", line 2963, in softmax
    return tf.nn.softmax(x, axis=axis)
TypeError: softmax() got an unexpected keyword argument 'axis'
就網上說是keras 2.1.6的softmax沒有axis這個參數了，於是，我回退了keras版本：

pip install keras==2.1
 然後程序就不報錯誤了，完工。