def step_HMC()

in example/bayesian-methods/algos.py [0:0]

• 40 lines of code
• 24 McCabe index (conditional complexity)

def step_HMC(exe, exe_params, exe_grads, label_key, noise_precision, prior_precision, L=10,
             eps=1E-6):
    init_params = {k: v.copyto(v.context) for k, v in exe_params.items()}
    end_params = {k: v.copyto(v.context) for k, v in exe_params.items()}
    init_momentums = {k: mx.random.normal(0, 1, v.shape) for k, v in init_params.items()}
    end_momentums = {k: v.copyto(v.context) for k, v in init_momentums.items()}
    init_potential = calc_potential(exe, init_params, label_key, noise_precision, prior_precision)

    # 0. Calculate Initial Energy and Kinetic
    init_kinetic = sum([nd.sum(nd.square(momentum)) / 2.0
                        for momentum in init_momentums.values()]).asscalar()
    # 1. Make a half step for momentum at the beginning
    exe.copy_params_from(end_params)
    exe.forward(is_train=True)
    exe.backward()
    for k, v in exe_grads.items():
        v.wait_to_read()
    for k, momentum in end_momentums.items():
        momentum[:] = momentum - (eps / 2) * exe_grads[k]
    # 2. Alternate full steps for position and momentum
    for i in range(L):
        # 2.1 Full step for position
        for k, param in exe_params.items():
            param[:] = param + eps * end_momentums[k]
        # 2.2 Full step for the momentum, except at the end of trajectory we perform a half step
        exe.forward(is_train=True)
        exe.backward()
        for v in exe_grads.values():
            v.wait_to_read()
        if i != L - 1:
            for k, momentum in end_momentums.items():
                momentum[:] = momentum - eps * exe_grads[k]
        else:
            for k, momentum in end_momentums.items():
                # We should reverse the sign of the momentum at the end
                momentum[:] = -(momentum - eps / 2.0 * exe_grads[k])
    copy_param(exe, end_params)
    # 3. Calculate acceptance ratio and accept/reject the move
    end_potential = calc_potential(exe, end_params, label_key, noise_precision, prior_precision)
    end_kinetic = sum([nd.sum(nd.square(momentum)) / 2.0
                       for momentum in end_momentums.values()]).asscalar()
    # print init_potential, init_kinetic, end_potential, end_kinetic
    r = numpy.random.rand(1)
    if r < numpy.exp(-(end_potential + end_kinetic) + (init_potential + init_kinetic)):
        exe.copy_params_from(end_params)
        return end_params, 1
    else:
        exe.copy_params_from(init_params)
        return init_params, 0