深度学习

tensorflow自定义梯度算法训练时间递增？

0

求助各位大神。。我用tensorflow自己写了一个momentum算法，问题是发现速度很慢很慢。。。实验室有一块titanX，跑tensorflow官网上的demo，超级快，2分钟就可以跑10多个epoch。。。但是同样的网络，只不过将train_step改成我自己定义的就非常慢，2个小时连mnist的1/3都跑不完。。。。

def momentum(loss, t, parameter_list):

    opt = GradientDescentOptimizer(1e-3)

    grads_and_vars = opt.compute_gradients(loss, parameter_list)

    capped_grads_and_vars = []

    for i in range(len(grads_and_vars)):

        gradient = grads_and_vars[i][0]

        variable = grads_and_vars[i][1]

        if t != 0:

            momentum = tf.add(gradient, tf.multiply(mu, last_momentum0[i]))

            capped_grads_and_vars.append((momentum, variable))

        else:

            momentum = gradient

            capped_grads_and_vars.append((momentum, variable))

    if t != 0:

        for i in range(len(grads_and_vars)):

            last_momentum0[i] = capped_grads_and_vars[i][0]

    else:

        for i in range(len(grads_and_vars)):

            last_momentum0.append(capped_grads_and_vars[i][0])

    return opt.apply_gradients(grads_and_vars)