def __init__(self, hiddens, layer_fn=nn.Linear, bias=True, norm_fn=None, activation_fn=nn.ReLU, dropout=None): super().__init__() # hack to prevent cycle imports from catalyst.modules.modules import name2nn layer_fn = name2nn(layer_fn) activation_fn = name2nn(activation_fn) norm_fn = name2nn(norm_fn) net = [] for i, (f_in, f_out) in enumerate(pairwise(hiddens)): net.append((f"layer_{i}", layer_fn(f_in, f_out, bias=bias))) if norm_fn is not None: net.append((f"norm_{i}", norm_fn(f_out))) if dropout is not None: net.append((f"drop_{i}", nn.Dropout(dropout))) if activation_fn is not None: net.append((f"activation_{i}", activation_fn())) self.net = torch.nn.Sequential(OrderedDict(net))
def __init__(self, state_shape, action_size, hiddens, layer_fn, activation_fn=nn.ReLU, norm_fn=None, bias=True, out_activation=nn.Tanh): super().__init__() # hack to prevent cycle imports from catalyst.modules.modules import name2nn layer_fn = name2nn(layer_fn) activation_fn = name2nn(activation_fn) norm_fn = name2nn(norm_fn) out_activation = name2nn(out_activation) state_size = reduce(lambda x, y: x * y, state_shape) self.feature_net = SequentialNet(hiddens=[state_size] + hiddens, layer_fn=layer_fn, activation_fn=activation_fn, norm_fn=norm_fn, bias=bias) self.policy_net = SequentialNet(hiddens=[hiddens[-1], action_size], layer_fn=nn.Linear, activation_fn=out_activation, norm_fn=None, bias=True) inner_init = create_optimal_inner_init(nonlinearity=activation_fn) self.feature_net.apply(inner_init) self.policy_net.apply(outer_init)
def __init__(self, arch="resnet50", pretrained=True, frozen=True, pooling=None, pooling_kwargs=None, embedding_size=None, bn_momentum=0.01, cut_layers=2): super().__init__() # hack to prevent cycle imports from catalyst.modules.modules import name2nn resnet = torchvision.models.__dict__[arch](pretrained=pretrained) modules = list(resnet.children())[:-cut_layers] # delete last layers if frozen: for param in modules: param.requires_grad = False if pooling is not None: pooling_kwargs = pooling_kwargs or {} pooling_layer = name2nn(pooling) pooling_fn = pooling_layer( in_features=resnet.fc.in_features, **pooling_kwargs) \ if "attn" in pooling.lower() \ else pooling_layer(**pooling_kwargs) modules += [pooling_fn] # @TODO: refactor if "concatattn" in pooling.lower(): resnet_out_features = resnet.fc.in_features * 3 elif any([ x in pooling.lower() for x in ["concat", "avgattn", "maxattn"] ]): resnet_out_features = resnet.fc.in_features * 2 else: resnet_out_features = resnet.fc.in_features else: resnet_out_features = resnet.fc.in_features flatten = name2nn("Flatten") modules += [flatten()] if embedding_size is not None: additional_modules = [ nn.Linear(resnet_out_features, embedding_size), nn.BatchNorm1d(num_features=embedding_size, momentum=bn_momentum) ] embeddings_weights_init(additional_modules) modules += additional_modules self.out_features = embedding_size else: self.out_features = resnet_out_features self.feature_net = nn.Sequential(*modules)
def __init__(self, state_shape, action_size, hiddens, layer_fn, concat_at=1, n_atoms=1, activation_fn=nn.ReLU, norm_fn=None, bias=True, out_activation=None): super().__init__() # hack to prevent cycle imports from catalyst.modules.modules import name2nn layer_fn = name2nn(layer_fn) activation_fn = name2nn(activation_fn) norm_fn = name2nn(norm_fn) out_activation = name2nn(out_activation) self.n_atoms = n_atoms state_size = reduce(lambda x, y: x * y, state_shape) if concat_at > 0: hiddens_ = [state_size] + hiddens[0:concat_at] self.observation_net = SequentialNet(hiddens=hiddens_, layer_fn=layer_fn, activation_fn=activation_fn, norm_fn=norm_fn, bias=bias) hiddens_ = \ [hiddens[concat_at - 1] + action_size] + hiddens[concat_at:] self.feature_net = SequentialNet(hiddens=hiddens_, layer_fn=layer_fn, activation_fn=activation_fn, norm_fn=norm_fn, bias=bias) else: self.observation_net = None hiddens_ = [state_size + action_size] + hiddens self.feature_net = SequentialNet(hiddens=hiddens_, layer_fn=layer_fn, activation_fn=activation_fn, norm_fn=norm_fn, bias=bias) self.value_net = SequentialNet(hiddens=[hiddens[-1], n_atoms], layer_fn=nn.Linear, activation_fn=out_activation, norm_fn=None, bias=True) inner_init = create_optimal_inner_init(nonlinearity=activation_fn) if self.observation_net is not None: self.observation_net.apply(inner_init) self.feature_net.apply(inner_init) self.value_net.apply(outer_init)
def __init__(self, state_shape, action_size, hiddens, layer_fn, activation_fn=nn.ReLU, norm_fn=None, bias=True, out_activation=nn.Sigmoid): super().__init__() # hack to prevent cycle imports from catalyst.modules.modules import name2nn self.n_action = action_size layer_fn = name2nn(layer_fn) activation_fn = name2nn(activation_fn) norm_fn = name2nn(norm_fn) out_activation = name2nn(out_activation) state_size = reduce(lambda x, y: x * y, state_shape) self.feature_net = SequentialNet(hiddens=[state_size] + hiddens, layer_fn=layer_fn, activation_fn=activation_fn, norm_fn=norm_fn, bias=bias) self.embedding_net = SequentialNet( hiddens=[hiddens[-1], action_size * 2], layer_fn=layer_fn, activation_fn=None, norm_fn=norm_fn, bias=bias) self.coupling1 = CouplingLayer(action_size=action_size, layer_fn=layer_fn, activation_fn=activation_fn, norm_fn=None, bias=bias, parity="odd") self.coupling2 = CouplingLayer(action_size=action_size, layer_fn=layer_fn, activation_fn=activation_fn, norm_fn=None, bias=bias, parity="even") self.squasher = SquashingLayer(out_activation) inner_init = create_optimal_inner_init(nonlinearity=activation_fn) self.feature_net.apply(inner_init) self.embedding_net.apply(inner_init)
def __init__(self, activation_fn=nn.Tanh): """ Layer that squashes samples from some distribution to be bounded. """ super().__init__() # hack to prevent cycle imports from catalyst.modules.modules import name2nn self.activation = name2nn(activation_fn)()
def __init__(self, state_shape, action_size, hiddens, layer_fn, activation_fn=nn.ReLU, norm_fn=None, bias=True, out_activation=nn.Tanh): super().__init__() # hack to prevent cycle imports from catalyst.modules.modules import name2nn layer_fn = name2nn(layer_fn) activation_fn = name2nn(activation_fn) norm_fn = name2nn(norm_fn) out_activation = name2nn(out_activation) state_size = state_shape[-1] self.feature_net = SequentialNet(hiddens=[state_size] + hiddens, layer_fn=layer_fn, activation_fn=activation_fn, norm_fn=norm_fn, bias=bias) self.attn = nn.Sequential( nn.Conv1d(in_channels=hiddens[-1], out_channels=1, kernel_size=1, bias=True), nn.Softmax(dim=1)) self.feature_net2 = SequentialNet( hiddens=[hiddens[-1] * 4, hiddens[-1]], layer_fn=layer_fn, activation_fn=activation_fn, norm_fn=norm_fn, bias=bias) self.policy_net = SequentialNet(hiddens=[hiddens[-1], action_size], layer_fn=nn.Linear, activation_fn=out_activation, norm_fn=None, bias=True) inner_init = create_optimal_inner_init(nonlinearity=activation_fn) self.feature_net.apply(inner_init) self.attn.apply(outer_init) self.feature_net2.apply(inner_init) self.policy_net.apply(outer_init)
def __init__(self, in_features, activation_fn="Softmax2d"): super().__init__() # hack to prevent cycle imports from catalyst.modules.modules import name2nn activation_fn = name2nn(activation_fn) self.attn = nn.Sequential( nn.Conv2d(in_features, 1, kernel_size=1, stride=1, padding=0, bias=False), activation_fn())
def __init__(self, arch="resnet50", pretrained=True, frozen=True, pooling="GlobalConcatPool2d", embedding_size=None, bn_momentum=0.01, cut_layers=1, **kwargs): super().__init__() # hack to prevent cycle imports from catalyst.modules.modules import name2nn resnet = torchvision.models.__dict__[arch](pretrained=pretrained) modules = list(resnet.children())[:-cut_layers] # delete last layers if frozen: for param in modules: param.requires_grad = False pooling = name2nn(pooling) flatten = name2nn("Flatten") modules += [ pooling(), flatten(), ] if embedding_size is not None: resnet_out_features = resnet.fc.in_features \ if "concat" not in pooling.lower() \ else resnet.fc.in_features * 2 additional_modules = [ nn.Linear(resnet_out_features, embedding_size), nn.BatchNorm1d(num_features=embedding_size, momentum=bn_momentum) ] embeddings_weights_init(additional_modules) modules += additional_modules self.feature_net = nn.Sequential(*modules)
def __init__( self, action_size, layer_fn, activation_fn=nn.ReLU, norm_fn=None, bias=True, parity="odd" ): """ Conditional affine coupling layer used in Real NVP Bijector. Original paper: https://arxiv.org/abs/1605.08803 Adaptation to RL: https://arxiv.org/abs/1804.02808 Important notes --------------- 1. State embeddings are supposed to have size (action_size * 2). 2. Scale and translation networks used in the Real NVP Bijector both have one hidden layer of (action_size) (activation_fn) units. 3. Parity ("odd" or "even") determines which part of the input is being copied and which is being transformed. """ super().__init__() # hack to prevent cycle imports from catalyst.modules.modules import name2nn layer_fn = name2nn(layer_fn) activation_fn = name2nn(activation_fn) norm_fn = name2nn(norm_fn) self.parity = parity if self.parity == "odd": self.copy_size = action_size // 2 else: self.copy_size = action_size - action_size // 2 self.scale_prenet = SequentialNet( hiddens=[action_size * 2 + self.copy_size, action_size], layer_fn=layer_fn, activation_fn=activation_fn, norm_fn=None, bias=bias ) self.scale_net = SequentialNet( hiddens=[action_size, action_size - self.copy_size], layer_fn=layer_fn, activation_fn=None, norm_fn=None, bias=True ) self.translation_prenet = SequentialNet( hiddens=[action_size * 2 + self.copy_size, action_size], layer_fn=layer_fn, activation_fn=activation_fn, norm_fn=None, bias=bias ) self.translation_net = SequentialNet( hiddens=[action_size, action_size - self.copy_size], layer_fn=layer_fn, activation_fn=None, norm_fn=None, bias=True ) inner_init = create_optimal_inner_init(nonlinearity=activation_fn) self.scale_prenet.apply(inner_init) self.scale_net.apply(outer_init) self.translation_prenet.apply(inner_init) self.translation_net.apply(outer_init)