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)
