def __init__(
self,
patch_size: int = 32,
kernel_type: str = 'polar', # 'cart' 'polar'
whitening: str = 'pcawt', # 'lw', 'pca', 'pcaws', 'pcawt
training_set: str = 'liberty', # 'liberty', 'notredame', 'yosemite'
output_dims: int = 128,
) -> None:
super().__init__()
relative: bool = kernel_type == 'polar'
sigma: float = 1.4 * (patch_size / 64)
# Sequence of modules.
smoothing = GaussianBlur2d((5, 5), (sigma, sigma), 'replicate')
gradients = MKDGradients()
ori = EmbedGradients(patch_size=patch_size, relative=relative)
ese = ExplicitSpacialEncoding(kernel_type=kernel_type,
fmap_size=patch_size,
in_dims=ori.kernel.d)
wh = Whitening(whitening,
load_whitening_model(kernel_type, training_set),
in_dims=ese.odims,
output_dims=output_dims)
self.features = nn.Sequential(smoothing, gradients, ori, ese, wh)
def monochrome_model(gaussian_p, solarize_p):
return nn.Sequential(
RandomResizedCrop((image_size, image_size),
interpolation="BICUBIC"), RandomHorizontalFlip(),
RandomApply(GaussianBlur2d(get_kernel_size(image_size),
(0.1, 2.0)),
p=gaussian_p), RandomSolarize(0, 0, p=solarize_p))
def color_model(gaussian_p, solarize_p):
return nn.Sequential(
RandomResizedCrop((image_size, image_size),
interpolation="BICUBIC"), RandomHorizontalFlip(),
ColorJitter(0.4, 0.4, 0.2, 0.1, p=0.8), RandomGrayscale(p=0.2),
RandomApply(GaussianBlur2d(get_kernel_size(image_size),
(0.1, 2.0)),
p=gaussian_p), RandomSolarize(0, 0, p=solarize_p))
def __init__(
self,
patch_size: int = 32,
kernel_type: str = 'concat',
whitening: str = 'pcawt',
training_set: str = 'liberty',
output_dims: int = 128,
) -> None:
super().__init__()
self.patch_size: int = patch_size
self.kernel_type: str = kernel_type
self.whitening: str = whitening
self.training_set: str = training_set
self.sigma = 1.4 * (patch_size / 64)
self.smoothing = GaussianBlur2d((5, 5), (self.sigma, self.sigma),
'replicate')
self.gradients = MKDGradients()
# This stupid thing needed for jitting...
polar_s: str = 'polar'
cart_s: str = 'cart'
self.parametrizations = [
polar_s, cart_s
] if self.kernel_type == 'concat' else [self.kernel_type]
# Initialize cartesian/polar embedding with absolute/relative gradients.
self.odims: int = 0
relative_orientations = {polar_s: True, cart_s: False}
self.feats = {}
for parametrization in self.parametrizations:
gradient_embedding = EmbedGradients(
patch_size=patch_size,
relative=relative_orientations[parametrization])
spatial_encoding = ExplicitSpacialEncoding(
kernel_type=parametrization,
fmap_size=patch_size,
in_dims=gradient_embedding.kernel.d)
self.feats[parametrization] = nn.Sequential(
gradient_embedding, spatial_encoding)
self.odims += spatial_encoding.odims
# Compute true output_dims.
self.output_dims: int = min(output_dims, self.odims)
# Load supervised(lw)/unsupervised(pca) model trained on training_set.
if self.whitening is not None:
whitening_models = torch.hub.load_state_dict_from_url(
urls[self.kernel_type],
map_location=lambda storage, loc: storage)
whitening_model = whitening_models[training_set]
self.whitening_layer = Whitening(whitening,
whitening_model,
in_dims=self.odims,
output_dims=self.output_dims)
self.odims = self.output_dims
self.eval()
def __init__(self, in_features, sigma=1, kernel_size=3):
super().__init__()
self.in_features = in_features
self.sigma = sigma
self.kernel_size = kernel_size
self.GaussianBlur2d = GaussianBlur2d(sigma=(self.sigma, self.sigma),
kernel_size=(self.kernel_size,
self.kernel_size),
border_type='constant')
def __init__(self):
super(Edge, self).__init__()
from kornia.filters import Laplacian, Sobel, GaussianBlur2d
layers = []
# layers.append(nn.Conv2d(1, 1, 3, padding = 1, bias=False))
# layers.append(nn.Conv2d(1, 1, 3, padding = 1, bias=False))
layers.append(GaussianBlur2d((5, 5), (1.5, 1.5)))
layers.append(Sobel(normalized=True))
# layers.append(Laplacian(7, normalized=True))
self.edge = nn.Sequential(*layers)
def __init__(self,
kernel_size: int or tuple = 3,
sigma: float or tuple = 1.,
border_type: str = 'reflect'):
super(MaskedGaussianBlur, self).__init__()
if isinstance(kernel_size, int):
self.kernel_size = (kernel_size, kernel_size)
else:
assert isinstance(kernel_size, tuple)
self.kernel_size = kernel_size
if isinstance(sigma, float):
self.sigma = (sigma, sigma)
else:
assert isinstance(sigma, tuple)
self.sigma = sigma
from kornia.filters import GaussianBlur2d
self.gblur = GaussianBlur2d(kernel_size=self.kernel_size,
sigma=(sigma, sigma),
border_type=border_type)
def __init__(self, input_shape, s=1.0, apply_transforms=None):
assert len(input_shape) == 3, "input_shape should be (H, W, C)"
self.input_shape = input_shape
self.H, self.W, self.C = input_shape[0], input_shape[1], input_shape[2]
self.s = s
self.apply_transforms = apply_transforms
if self.apply_transforms is None:
kernel_size = int(0.1 * self.H)
sigma = self._get_sigma()
self.apply_transforms = KorniaCompose([
RandomResizedCrop(size=(self.H, self.W), scale=(0.08, 1.0)),
RandomHorizontalFlip(p=0.5),
ColorJitter(0.8 * self.s, 0.8 * self.s, 0.8 * self.s, 0.2 * self.s),
RandomGrayscale(p=0.2),
GaussianBlur2d(kernel_size=(kernel_size, kernel_size),
sigma=(sigma, sigma))
])
def gaussian_blur(self, p: float = 1.0) -> TransformType:
return RandomApply(GaussianBlur2d((3, 3), (1.5, 1.5)), p)
def __init__(
self,
image_shape,
output_size,
n_atoms,
dueling,
jumps,
spr,
augmentation,
target_augmentation,
eval_augmentation,
dynamics_blocks,
norm_type,
noisy_nets,
aug_prob,
classifier,
imagesize,
time_offset,
local_spr,
global_spr,
momentum_encoder,
shared_encoder,
distributional,
dqn_hidden_size,
momentum_tau,
renormalize,
renormalize_type,
q_l1_type,
dropout,
final_classifier,
model_rl,
noisy_nets_std,
residual_tm,
pred_hidden_ratio,
encoder_type,
transition_type,
conv_proj_channel,
proj_hidden_size,
gru_input_size,
gru_proj_size,
ln_ratio,
use_maxpool=False,
channels=None, # None uses default.
kernel_sizes=None,
strides=None,
paddings=None,
framestack=4,
):
"""Instantiates the neural network according to arguments; network defaults
stored within this method."""
super().__init__()
self.noisy = noisy_nets
self.time_offset = time_offset
self.aug_prob = aug_prob
self.classifier_type = classifier
self.distributional = distributional
n_atoms = 1 if not self.distributional else n_atoms
self.dqn_hidden_size = dqn_hidden_size
self.transforms = []
self.eval_transforms = []
self.uses_augmentation = False
for aug in augmentation:
if aug == "affine":
transformation = RandomAffine(5, (.14, .14), (.9, 1.1), (-5, 5))
eval_transformation = nn.Identity()
self.uses_augmentation = True
elif aug == "crop":
transformation = RandomCrop((84, 84))
# Crashes if aug-prob not 1: use CenterCrop((84, 84)) or Resize((84, 84)) in that case.
eval_transformation = CenterCrop((84, 84))
self.uses_augmentation = True
imagesize = 84
elif aug == "rrc":
transformation = RandomResizedCrop((100, 100), (0.8, 1))
eval_transformation = nn.Identity()
self.uses_augmentation = True
elif aug == "blur":
transformation = GaussianBlur2d((5, 5), (1.5, 1.5))
eval_transformation = nn.Identity()
self.uses_augmentation = True
elif aug == "shift":
transformation = nn.Sequential(nn.ReplicationPad2d(4), RandomCrop((84, 84)))
eval_transformation = nn.Identity()
elif aug == "intensity":
transformation = Intensity(scale=0.05)
eval_transformation = nn.Identity()
elif aug == "none":
transformation = eval_transformation = nn.Identity()
else:
raise NotImplementedError()
self.transforms.append(transformation)
self.eval_transforms.append(eval_transformation)
self.dueling = dueling
f, c = image_shape[:2]
in_channels = np.prod(image_shape[:2])
if encoder_type == 'conv2d':
self.conv = Conv2dModel(
in_channels=in_channels,
channels=[32, 64, 64],
kernel_sizes=[8, 4, 3],
strides=[4, 2, 1],
paddings=[0, 0, 0],
use_maxpool=False,
dropout=dropout,
conv_proj_channel=conv_proj_channel,
)
elif encoder_type == 'resnet18':
self.conv = resnet18()
else:
raise NotImplementedError
fake_input = torch.zeros(1, f*c, imagesize, imagesize)
fake_output = self.conv(fake_input)
self.hidden_size = fake_output.shape[1]
self.pixels = fake_output.shape[-1]*fake_output.shape[-2]
print("Spatial latent size is {}".format(fake_output.shape[1:]))
if proj_hidden_size:
self.conv_proj = nn.Sequential(
nn.Flatten(1, -1),
nn.Linear(self.hidden_size * self.pixels, proj_hidden_size),
nn.LayerNorm(proj_hidden_size),
nn.ReLU(),
nn.Dropout(dropout),
)
else:
self.conv_proj = nn.Identity()
self.jumps = jumps
self.model_rl = model_rl
self.use_spr = spr
self.target_augmentation = target_augmentation
self.eval_augmentation = eval_augmentation
self.num_actions = output_size
self.transition_type = transition_type
if dueling:
self.head = DQNDistributionalDuelingHeadModel(self.hidden_size,
output_size,
hidden_size=self.dqn_hidden_size,
pixels=self.pixels,
noisy=self.noisy,
n_atoms=n_atoms,
std_init=noisy_nets_std,
proj_hidden_size=proj_hidden_size)
else:
self.head = DQNDistributionalHeadModel(self.hidden_size,
output_size,
hidden_size=self.dqn_hidden_size,
pixels=self.pixels,
noisy=self.noisy,
n_atoms=n_atoms,
std_init=noisy_nets_std)
if self.jumps > 0:
repr_size = proj_hidden_size if proj_hidden_size else (self.pixels * self.hidden_size)
if transition_type == 'gru':
self.dynamics_model = GRUModel(
input_size = gru_input_size,
repr_size = repr_size,
proj_size = gru_proj_size,
num_layers = 1,
num_actions = self.num_actions,
renormalize=renormalize,
renormalize_type=renormalize_type,
dropout=dropout
)
else:
self.dynamics_model = TransitionModel(channels=self.hidden_size,
num_actions=output_size,
pixels=self.pixels,
hidden_size=self.hidden_size,
limit=1,
blocks=dynamics_blocks,
norm_type=norm_type,
renormalize=renormalize,
residual=residual_tm)
else:
self.dynamics_model = nn.Identity()
self.renormalize = renormalize
self.renormalize_type = renormalize_type
self.ln_ratio = ln_ratio
if renormalize_type == 'train_ln':
self.renormalize_ln = nn.LayerNorm(repr_size)
else:
self.renormalize_ln = nn.Identity()
if self.use_spr:
self.local_spr = local_spr
self.global_spr = global_spr
self.momentum_encoder = momentum_encoder
self.momentum_tau = momentum_tau
self.shared_encoder = shared_encoder
assert not (self.shared_encoder and self.momentum_encoder)
# in case someone tries something silly like --local-spr 2
self.num_sprs = int(bool(self.local_spr)) + \
int(bool(self.global_spr))
if self.local_spr:
self.local_final_classifier = nn.Identity()
if self.classifier_type == "mlp":
self.local_classifier = nn.Sequential(nn.Linear(self.hidden_size,
self.hidden_size),
nn.BatchNorm1d(self.hidden_size),
nn.ReLU(),
nn.Linear(self.hidden_size,
self.hidden_size))
elif self.classifier_type == "bilinear":
self.local_classifier = nn.Linear(self.hidden_size, self.hidden_size)
elif self.classifier_type == "none":
self.local_classifier = nn.Identity()
if final_classifier == "mlp":
self.local_final_classifier = nn.Sequential(nn.Linear(self.hidden_size, 2*self.hidden_size),
nn.BatchNorm1d(2*self.hidden_size),
nn.ReLU(),
nn.Linear(2*self.hidden_size,
self.hidden_size))
elif final_classifier == "linear":
self.local_final_classifier = nn.Linear(self.hidden_size, self.hidden_size)
else:
self.local_final_classifier = nn.Identity()
self.local_target_classifier = self.local_classifier
else:
self.local_classifier = self.local_target_classifier = nn.Identity()
if self.global_spr:
self.global_final_classifier = nn.Identity()
if self.classifier_type == "mlp":
self.global_classifier = nn.Sequential(
nn.Flatten(-3, -1),
nn.Linear(self.pixels*self.hidden_size, 512),
nn.BatchNorm1d(512),
nn.ReLU(),
nn.Linear(512, 256)
)
self.global_target_classifier = self.global_classifier
global_spr_size = 256
elif self.classifier_type == "q_l1":
self.global_classifier = QL1Head(self.head, dueling=dueling, type=q_l1_type)
global_spr_size = self.global_classifier.out_features
self.global_target_classifier = self.global_classifier
elif self.classifier_type == "q_l2":
self.global_classifier = nn.Sequential(self.head, nn.Flatten(-2, -1))
self.global_target_classifier = self.global_classifier
global_spr_size = 256
elif self.classifier_type == "bilinear":
self.global_classifier = nn.Sequential(nn.Flatten(-3, -1),
nn.Linear(self.hidden_size*self.pixels,
self.hidden_size*self.pixels))
self.global_target_classifier = nn.Flatten(-3, -1)
elif self.classifier_type == "none":
self.global_classifier = nn.Flatten(-3, -1)
self.global_target_classifier = nn.Flatten(-3, -1)
global_spr_size = self.hidden_size*self.pixels
if final_classifier == "mlp":
global_final_hidden_size = int(global_spr_size * pred_hidden_ratio)
self.global_final_classifier = nn.Sequential(
nn.Linear(global_spr_size, global_final_hidden_size),
nn.BatchNorm1d(global_final_hidden_size),
nn.ReLU(),
nn.Linear(global_final_hidden_size, global_spr_size)
)
elif final_classifier == "linear":
self.global_final_classifier = nn.Sequential(
nn.Linear(global_spr_size, global_spr_size),
)
elif final_classifier == "none":
self.global_final_classifier = nn.Identity()
else:
self.global_classifier = self.global_target_classifier = nn.Identity()
if self.momentum_encoder:
self.target_encoder = copy.deepcopy(self.conv)
self.target_encoder_proj = copy.deepcopy(self.conv_proj)
self.target_renormalize_ln = copy.deepcopy(self.renormalize_ln)
self.global_target_classifier = copy.deepcopy(self.global_target_classifier)
self.local_target_classifier = copy.deepcopy(self.local_target_classifier)
for param in (list(self.target_encoder.parameters())
+ list(self.target_encoder_proj.parameters())
+ list(self.target_renormalize_ln.parameters())
+ list(self.global_target_classifier.parameters())
+ list(self.local_target_classifier.parameters())):
param.requires_grad = False
elif not self.shared_encoder:
# Use a separate target encoder on the last frame only.
self.global_target_classifier = copy.deepcopy(self.global_target_classifier)
self.local_target_classifier = copy.deepcopy(self.local_target_classifier)
if self.stack_actions:
input_size = c - 1
else:
input_size = c
self.target_encoder = Conv2dModel(in_channels=input_size,
channels=[32, 64, 64],
kernel_sizes=[8, 4, 3],
strides=[4, 2, 1],
paddings=[0, 0, 0],
use_maxpool=False,
)
elif self.shared_encoder:
self.target_encoder = self.conv
print("Initialized model with {} parameters".format(count_parameters(self)))
