mlp.py

MLP 编码基础

原文件 · 56 行 · 1947 字节

algorithms/utils/mlp.py
SHA-256: a0475371f8454bc02111d8a6e27916e207a01a6c45eba8b4a52a83d9bfd65e5a

1import torch.nn as nn2from .util import init, get_clones34"""MLP modules."""56class MLPLayer(nn.Module):7    def __init__(self, input_dim, hidden_size, layer_N, use_orthogonal, use_ReLU):8        super(MLPLayer, self).__init__()9        self._layer_N = layer_N1011        active_func = [nn.Tanh(), nn.ReLU()][use_ReLU]12        init_method = [nn.init.xavier_uniform_, nn.init.orthogonal_][use_orthogonal]13        gain = nn.init.calculate_gain(['tanh', 'relu'][use_ReLU])1415        def init_(m):16            return init(m, init_method, lambda x: nn.init.constant_(x, 0), gain=gain)1718        self.fc1 = nn.Sequential(19            init_(nn.Linear(input_dim, hidden_size)), active_func, nn.LayerNorm(hidden_size))20        self.fc_h = nn.Sequential(init_(21            nn.Linear(hidden_size, hidden_size)), active_func, nn.LayerNorm(hidden_size))22        self.fc2 = get_clones(self.fc_h, self._layer_N)2324    def forward(self, x):25        x = self.fc1(x)26        for i in range(self._layer_N):27            x = self.fc2[i](x)28        return x293031class MLPBase(nn.Module):32    def __init__(self, args, obs_shape, cat_self=True, attn_internal=False):33        super(MLPBase, self).__init__()3435        self._use_feature_normalization = args.use_feature_normalization36        self._use_orthogonal = args.use_orthogonal37        self._use_ReLU = args.use_ReLU38        self._stacked_frames = args.stacked_frames39        self._layer_N = args.layer_N40        self.hidden_size = args.hidden_size4142        obs_dim = obs_shape[0]4344        if self._use_feature_normalization:45            self.feature_norm = nn.LayerNorm(obs_dim)4647        self.mlp = MLPLayer(obs_dim, self.hidden_size,48                              self._layer_N, self._use_orthogonal, self._use_ReLU)4950    def forward(self, x):51        if self._use_feature_normalization:52            x = self.feature_norm(x)5354        x = self.mlp(x)5556        return x