train.py 文件源码

def get_loss(logit, label):
    label = tf.to_int64(label)
    cross_entropy = tf.nn.sparse_softmax_cross_entropy_with_logits(
                                                        logits=logit,
                                                        labels=label,
                                                        name='cross_entropy')

    loss = tf.reduce_mean(cross_entropy)

    return loss

## get optimizer
# @param learning_rate:
# @param optimizer: optimizer method