GLPhysXRL/RunningMean.cpp at main · MihaiAnca13/GLPhysXRL · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
//
// Created by mihai on 29/04/23.
//

#include "RunningMean.h"
#include <iostream>
using namespace std;


RunningMeanStdImpl::RunningMeanStdImpl(IntArrayRef insize, double epsilon, bool per_channel, bool norm_only)
        : insize(insize), epsilon(epsilon), norm_only(norm_only), per_channel(per_channel) {
    if (per_channel) {
        if (insize.size() == 3) {
            axis = {0, 2, 3};
        }
        if (insize.size() == 2) {
            axis = {0, 2};
        }
        if (insize.size() == 1) {
            axis = {0};
        }
        in_size = insize[0];
    } else {
        axis = {0};
        in_size = insize.size();
    }

    running_mean = register_buffer("running_mean", torch::zeros({in_size}, kFloat64));
    running_var = register_buffer("running_var", torch::ones({in_size}, kFloat64));
    count = register_buffer("count", torch::ones({}, kFloat64));
}

void RunningMeanStdImpl::_update_mean_var_count_from_moments(const Tensor& batch_mean, const Tensor& batch_var, const Tensor& batch_count) {
    Tensor delta = batch_mean - running_mean;
    Tensor tot_count = count + batch_count;

    running_mean = running_mean + delta * batch_count / tot_count;

    Tensor m_a = running_var * count;
    Tensor m_b = batch_var * batch_count;
    Tensor M2 = m_a + m_b + delta.pow(2) * count * batch_count / tot_count;

    running_var = M2 / tot_count;
    count = tot_count;
}

Tensor RunningMeanStdImpl::forward(Tensor input, bool unnorm) {
    if (is_training()) {
        auto mean = input.mean(axis);
        auto var = input.var(axis);

        _update_mean_var_count_from_moments(mean, var, torch::tensor(input.size(0)));
    }

    Tensor current_mean, current_var;

    if (per_channel) {
        if (insize.size() == 3) {
            current_mean = running_mean.view({1, insize[0], 1, 1}).expand_as(input);
            current_var = running_var.view({1, insize[0], 1, 1}).expand_as(input);
        }
        if (insize.size() == 2) {
            current_mean = running_mean.view({1, insize[0], 1}).expand_as(input);
            current_var = running_var.view({1, insize[0], 1}).expand_as(input);
        }
        if (insize.size() == 1) {
            current_mean = running_mean.view({1, insize[0]}).expand_as(input);
            current_var = running_var.view({1, insize[0]}).expand_as(input);
        }
    } else {
        current_mean = running_mean;
        current_var = running_var;
    }

    Tensor y;
    if (unnorm) {
        y = input.clamp(-5.0, 5.0);
        y = sqrt(current_var.to(kFloat32) + epsilon) * y.to(kFloat32) + current_mean.to(kFloat32);
    } else {
        if (norm_only) {
            y = input / sqrt(current_var.to(kFloat32) + epsilon);
        } else {
            y = (input - current_mean.to(kFloat32)) / sqrt(current_var.to(kFloat32) + epsilon);
            y = y.clamp(-5.0, 5.0);
        }
    }
    return y;
}