Tensorflow Keras 交叉熵损失与 output 上的线性激活

Question

在 PyTorch 中，交叉熵损失 function 类似于

CrossEntropyLoss(x, y) = H(one_hot(y), softmax(x))

所以你可以有一个线性 output 层。 有没有办法用 tf.keras.Sequential 做到这一点？
我为 MNIST 写了这个小 CNN

model = tf.keras.Sequential()
model.add(tfkl.Input(shape=(28, 28, 1)))
model.add(tfkl.Conv2D(32, (5, 5), padding="valid", activation=tf.nn.relu))
model.add(tfkl.MaxPool2D((2, 2)))
model.add(tfkl.Conv2D(64, (5, 5), padding="valid", activation=tf.nn.relu))
model.add(tfkl.MaxPool2D((2, 2)))
model.add(tfkl.Flatten())
model.add(tfkl.Dense(1024, activation=tf.nn.relu))
model.add(tfkl.Dense(10, activation=tf.nn.softmax))


model.compile(optimizer="adam",
              loss="sparse_categorical_crossentropy",
              metrics=["accuracy"])
model.summary()
model.fit(x_train, y_train, epochs=1)

我想拥有

model.add(tfkl.Dense(10))

作为最后一层。
我正在尝试实现 ADef 算法，但梯度 wrt 的条目。 输入似乎太小了，我猜用线性 output 他们会是对的。
我知道有 tf.nn.softmax_cross_entropy_with_logits 但我不知道如何在这种情况下使用它。

编辑：
改变

model.compile(optimizer="adam",
              loss="sparse_categorical_crossentropy",
              metrics=["accuracy"])

至

model.compile(optimizer="adam",
              loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
              metrics=["accuracy"])

成功了。

Answer 1

谢谢@Moe1234。 为了社区的利益，在这里提供解决方案

更改后问题已解决

model.compile(optimizer="adam",
              loss="sparse_categorical_crossentropy",
              metrics=["accuracy"])

至

model.compile(optimizer="adam",
              loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
              metrics=["accuracy"])