net_dropped = torch.nn.Sequential(

    torch.nn.Linear(1, N_HIDDEN),

    torch.nn.Dropout(0.5),  # drop 50% of the neuron

    torch.nn.ReLU(),

    torch.nn.Linear(N_HIDDEN, N_HIDDEN),

    torch.nn.Dropout(0.5),  # drop 50% of the neuron

    torch.nn.ReLU(),

    torch.nn.Linear(N_HIDDEN, 1),

)



for t in range(500):

    pred_drop = net_dropped(x)

    loss_drop = loss_func(pred_drop, y)

    optimizer_drop.zero_grad()

    loss_drop.backward()

    optimizer_drop.step()

    if t % 10 == 0:

        # change to eval mode in order to fix drop out effect

        net_dropped.eval()  # parameters for dropout differ from train mode

        test_pred_drop = net_dropped(test_x)

        # change back to train mode

        net_dropped.train()

            if self.do_bn:

                bn = nn.BatchNorm1d(10, momentum=0.5)

                setattr(self, 'bn%i' % i, bn)   # IMPORTANT set layer to the Module

                self.bns.append(bn)

    for epoch in range(EPOCH):

        print('Epoch: ', epoch)

        for net, l in zip(nets, losses):

            net.eval()              # set eval mode to fix moving_mean and moving_var

            pred, layer_input, pre_act = net(test_x)

            net.train()             # free moving_mean and moving_var

        plot_histogram(*layer_inputs, *pre_acts)  

tf_is_training = tf.placeholder(tf.bool, None)  # to control dropout when training and testing

# dropout net

d1 = tf.layers.dense(tf_x, N_HIDDEN, tf.nn.relu)

d1 = tf.layers.dropout(d1, rate=0.5, training=tf_is_training)   # drop out 50% of inputs


d2 = tf.layers.dense(d1, N_HIDDEN, tf.nn.relu)

d2 = tf.layers.dropout(d2, rate=0.5, training=tf_is_training)   # drop out 50% of inputs


d_out = tf.layers.dense(d2, 1)

for t in range(500):

    sess.run([o_train, d_train], {tf_x: x, tf_y: y, tf_is_training: True})  # train, set is_training=True

    if t % 10 == 0:

        # plotting

        plt.cla()

        o_loss_, d_loss_, o_out_, d_out_ = sess.run(

            [o_loss, d_loss, o_out, d_out], {tf_x: test_x, tf_y: test_y, tf_is_training: False} # test, set is_training=False

        )

    def add_layer(self, x, out_size, ac=None):

        x = tf.layers.dense(x, out_size, kernel_initializer=self.w_init, bias_initializer=B_INIT)

        self.pre_activation.append(x)

        # the momentum plays important rule. the default 0.99 is too high in this case!

        if self.is_bn: x = tf.layers.batch_normalization(x, momentum=0.4, training=tf_is_train)    # when have BN

        out = x if ac is None else ac(x)

        return out

        # !! IMPORTANT !! the moving_mean and moving_variance need to be updated,

        # pass the update_ops with control_dependencies to the train_op

        update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)

        with tf.control_dependencies(update_ops):

            self.train = tf.train.AdamOptimizer(LR).minimize(self.loss)