def __init__(self,
action_size,
layer_fn,
activation_fn=nn.ReLU,
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__()
layer_fn = MODULES.get_if_str(layer_fn)
activation_fn = MODULES.get_if_str(activation_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)
def __init__(self,
hiddens,
layer_fn=nn.Linear,
bias=True,
norm_fn=None,
activation_fn=nn.ReLU,
dropout=None,
layer_order=None,
residual=False):
super().__init__()
assert len(hiddens) > 1, "No sequence found"
layer_fn = MODULES.get_if_str(layer_fn)
activation_fn = MODULES.get_if_str(activation_fn)
norm_fn = MODULES.get_if_str(norm_fn)
dropout = MODULES.get_if_str(dropout)
inner_init = create_optimal_inner_init(nonlinearity=activation_fn)
layer_order = layer_order or ["layer", "norm", "drop", "act"]
if isinstance(dropout, float):
dropout_fn = lambda: nn.Dropout(dropout)
else:
dropout_fn = dropout
def _layer_fn(f_in, f_out, bias):
return layer_fn(f_in, f_out, bias=bias)
def _normalize_fn(f_in, f_out, bias):
return norm_fn(f_out) if norm_fn is not None else None
def _dropout_fn(f_in, f_out, bias):
return dropout_fn() if dropout_fn is not None else None
def _activation_fn(f_in, f_out, bias):
return activation_fn() if activation_fn is not None else None
name2fn = {
"layer": _layer_fn,
"norm": _normalize_fn,
"drop": _dropout_fn,
"act": _activation_fn,
}
net = []
for i, (f_in, f_out) in enumerate(pairwise(hiddens)):
block = []
for key in layer_order:
fn = name2fn[key](f_in, f_out, bias)
if fn is not None:
block.append((f"{key}", fn))
block = torch.nn.Sequential(OrderedDict(block))
if residual:
block = ResidualWrapper(net=block)
net.append((f"block_{i}", block))
self.net = torch.nn.Sequential(OrderedDict(net))
self.net.apply(inner_init)
def _normalization_fn(normalization_fn, f_in, f_out, **kwargs):
normalization_fn = MODULES.get_if_str(normalization_fn)
normalization_fn = \
normalization_fn(f_out, **kwargs) \
if normalization_fn is not None \
else None
return normalization_fn
def __init__(
self,
input_size=224,
width_mult=1.,
pretrained=True,
pooling=None,
pooling_kwargs=None,
):
super().__init__()
net = MobileNetV2(input_size=input_size,
width_mult=width_mult,
pretrained=pretrained)
self.encoder = list(net.encoder.children())
if pooling is not None:
pooling_kwargs = pooling_kwargs or {}
pooling_layer_fn = MODULES.get(pooling)
pooling_layer = pooling_layer_fn(
in_features=self.last_channel, **pooling_kwargs) \
if "attn" in pooling.lower() \
else pooling_layer_fn(**pooling_kwargs)
self.encoder.append(pooling_layer)
out_features = pooling_layer.out_features(
in_features=net.output_channel)
else:
out_features = net.output_channel
self.out_features = out_features
# make it nn.Sequential
self.encoder = nn.Sequential(*self.encoder)
self._initialize_weights()
def __init__(self, squashing_fn=nn.Tanh):
"""
Layer that squashes samples from some distribution to be bounded.
"""
super().__init__()
self.squashing_fn = MODULES.get_if_str(squashing_fn)()
def __init__(self,
arch="resnet34",
pretrained=True,
frozen=True,
pooling=None,
pooling_kwargs=None,
cut_layers=2):
super().__init__()
resnet = torchvision.models.__dict__[arch](pretrained=pretrained)
modules = list(resnet.children())[:-cut_layers] # delete last layers
if frozen:
for module in modules:
for param in module.parameters():
param.requires_grad = False
if pooling is not None:
pooling_kwargs = pooling_kwargs or {}
pooling_layer_fn = MODULES.get(pooling)
pooling_layer = pooling_layer_fn(
in_features=resnet.fc.in_features, **pooling_kwargs) \
if "attn" in pooling.lower() \
else pooling_layer_fn(**pooling_kwargs)
modules += [pooling_layer]
out_features = pooling_layer.out_features(
in_features=resnet.fc.in_features)
else:
out_features = resnet.fc.in_features
modules += [Flatten()]
self.out_features = out_features
self.encoder = nn.Sequential(*modules)
def _get_block(**conv_params):
layers = [nn.Conv2d(**conv_params)]
if normalization is not None:
normalization_fn = MODULES.get_if_str(normalization)
layers.append(normalization_fn(conv_params["out_channels"]))
if dropout_rate is not None:
layers.append(nn.Dropout2d(p=dropout_rate))
layers.append(activation_fn(inplace=True))
return layers
def __init__(self, in_features, activation_fn="Sigmoid"):
super().__init__()
activation_fn = MODULES.get_if_str(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: str = "resnet18",
pretrained: bool = True,
frozen: bool = True,
pooling: str = None,
pooling_kwargs: dict = None,
cut_layers: int = 2,
state_dict: Union[dict, str, Path] = None,
):
"""
Args:
arch (str): Name for resnet. Have to be one of
resnet18, resnet34, resnet50, resnet101, resnet152
pretrained (bool): If True, returns a model pre-trained on ImageNet
frozen (bool): If frozen, sets requires_grad to False
pooling (str): pooling
pooling_kwargs (dict): params for pooling
state_dict (Union[dict, str, Path]): Path to ``torch.Model``
or a dict containing parameters and persistent buffers.
"""
super().__init__()
resnet = torchvision.models.__dict__[arch](pretrained=pretrained)
if state_dict is not None:
if isinstance(state_dict, (Path, str)):
state_dict = torch.load(str(state_dict))
resnet.load_state_dict(state_dict)
modules = list(resnet.children())[:-cut_layers] # delete last layers
if frozen:
for module in modules:
utils.set_requires_grad(module, requires_grad=False)
if pooling is not None:
pooling_kwargs = pooling_kwargs or {}
pooling_layer_fn = MODULES.get(pooling)
pooling_layer = (pooling_layer_fn(
in_features=resnet.fc.in_features, **pooling_kwargs) if "attn"
in pooling.lower() else pooling_layer_fn(
**pooling_kwargs))
modules += [pooling_layer]
if hasattr(pooling_layer, "out_features"):
out_features = pooling_layer.out_features(
in_features=resnet.fc.in_features)
else:
out_features = None
else:
out_features = resnet.fc.in_features
modules += [Flatten()]
self.out_features = out_features
self.encoder = nn.Sequential(*modules)
def __init__(self,
action_size,
layer_fn,
activation_fn=nn.ReLU,
squashing_fn=nn.Tanh,
bias=False):
super().__init__()
activation_fn = MODULES.get_if_str(activation_fn)
self.action_size = action_size
self.coupling1 = CouplingLayer(action_size=action_size,
layer_fn=layer_fn,
activation_fn=activation_fn,
bias=bias,
parity="odd")
self.coupling2 = CouplingLayer(action_size=action_size,
layer_fn=layer_fn,
activation_fn=activation_fn,
bias=bias,
parity="even")
self.squashing_layer = SquashingLayer(squashing_fn)
def _activation_fn(activation_fn, f_in, f_out, **kwargs):
activation_fn = MODULES.get_if_str(activation_fn)
activation_fn = activation_fn(**kwargs) \
if activation_fn is not None \
else None
return activation_fn
def _dropout_fn(dropout_fn, f_in, f_out, **kwargs):
dropout_fn = MODULES.get_if_str(dropout_fn)
dropout_fn = dropout_fn(**kwargs) \
if dropout_fn is not None \
else None
return dropout_fn
def _layer_fn(layer_fn, f_in, f_out, **kwargs):
layer_fn = MODULES.get_if_str(layer_fn)
layer_fn = layer_fn(f_in, f_out, **kwargs)
return layer_fn