valuenorm.py

价值归一化

原文件 · 79 行 · 3189 字节

utils/valuenorm.py
SHA-256: 8cd48bb7e3a3d8295baec57a2cde23ca60a973d047f358f408709ef53c6a1836

12import numpy as np34import torch5import torch.nn as nn678class ValueNorm(nn.Module):9    """ Normalize a vector of observations - across the first norm_axes dimensions"""1011    def __init__(self, input_shape, norm_axes=1, beta=0.99999, per_element_update=False, epsilon=1e-5, device=torch.device("cpu")):12        super(ValueNorm, self).__init__()1314        self.input_shape = input_shape15        self.norm_axes = norm_axes16        self.epsilon = epsilon17        self.beta = beta18        self.per_element_update = per_element_update19        self.tpdv = dict(dtype=torch.float32, device=device)2021        self.running_mean = nn.Parameter(torch.zeros(input_shape), requires_grad=False).to(**self.tpdv)22        self.running_mean_sq = nn.Parameter(torch.zeros(input_shape), requires_grad=False).to(**self.tpdv)23        self.debiasing_term = nn.Parameter(torch.tensor(0.0), requires_grad=False).to(**self.tpdv)24        25        self.reset_parameters()2627    def reset_parameters(self):28        self.running_mean.zero_()29        self.running_mean_sq.zero_()30        self.debiasing_term.zero_()3132    def running_mean_var(self):33        debiased_mean = self.running_mean / self.debiasing_term.clamp(min=self.epsilon)34        debiased_mean_sq = self.running_mean_sq / self.debiasing_term.clamp(min=self.epsilon)35        debiased_var = (debiased_mean_sq - debiased_mean ** 2).clamp(min=1e-2)36        return debiased_mean, debiased_var3738    @torch.no_grad()39    def update(self, input_vector):40        if type(input_vector) == np.ndarray:41            input_vector = torch.from_numpy(input_vector)42        input_vector = input_vector.to(**self.tpdv)4344        batch_mean = input_vector.mean(dim=tuple(range(self.norm_axes)))45        batch_sq_mean = (input_vector ** 2).mean(dim=tuple(range(self.norm_axes)))4647        if self.per_element_update:48            batch_size = np.prod(input_vector.size()[:self.norm_axes])49            weight = self.beta ** batch_size50        else:51            weight = self.beta5253        self.running_mean.mul_(weight).add_(batch_mean * (1.0 - weight))54        self.running_mean_sq.mul_(weight).add_(batch_sq_mean * (1.0 - weight))55        self.debiasing_term.mul_(weight).add_(1.0 * (1.0 - weight))5657    def normalize(self, input_vector):58        # Make sure input is float3259        if type(input_vector) == np.ndarray:60            input_vector = torch.from_numpy(input_vector)61        input_vector = input_vector.to(**self.tpdv)6263        mean, var = self.running_mean_var()64        out = (input_vector - mean[(None,) * self.norm_axes]) / torch.sqrt(var)[(None,) * self.norm_axes]65        66        return out6768    def denormalize(self, input_vector):69        """ Transform normalized data back into original distribution """70        if type(input_vector) == np.ndarray:71            input_vector = torch.from_numpy(input_vector)72        input_vector = input_vector.to(**self.tpdv)7374        mean, var = self.running_mean_var()75        out = input_vector * torch.sqrt(var)[(None,) * self.norm_axes] + mean[(None,) * self.norm_axes]76        77        out = out.cpu().numpy()78        79        return out