def initialize_torchvision_model(name, d_out, **kwargs):
import torchvision
# get constructor and last layer names
if name == 'wideresnet50':
constructor_name = 'wide_resnet50_2'
last_layer_name = 'fc'
elif name == 'densenet121':
constructor_name = name
last_layer_name = 'classifier'
elif name in ('resnet18', 'resnet34', 'resnet50', 'resnet101'):
constructor_name = name
last_layer_name = 'fc'
else:
raise ValueError(f'Torchvision model {name} not recognized')
# construct the default model, which has the default last layer
constructor = getattr(torchvision.models, constructor_name)
model = constructor(**kwargs)
# adjust the last layer
d_features = getattr(model, last_layer_name).in_features
if d_out is None: # want to initialize a featurizer model
last_layer = Identity(d_features)
model.d_out = d_features
else: # want to initialize a classifier for a particular num_classes
last_layer = nn.Linear(d_features, d_out)
model.d_out = d_out
setattr(model, last_layer_name, last_layer)
return model
def __init__(
self,
input_height=28,
input_channels=1,
noise_dim=100,
z_dim=32,
nonlinearity='softplus',
enc_noise=False,
):
super().__init__()
self.input_height = input_height
self.input_channels = input_channels
self.noise_dim = noise_dim
self.z_dim = z_dim
self.nonlinearity = nonlinearity
self.enc_noise = enc_noise
h_dim = 256
nos_dim = noise_dim if not enc_noise else h_dim
s_h = input_height
s_h2 = conv_out_size(s_h, 5, 2, 2)
s_h4 = conv_out_size(s_h2, 5, 2, 2)
s_h8 = conv_out_size(s_h4, 5, 2, 2)
#print(s_h, s_h2, s_h4, s_h8)
self.afun = get_nonlinear_func(nonlinearity)
self.conv1 = nn.Conv2d(self.input_channels, 16, 5, 2, 2, bias=True)
self.conv2 = nn.Conv2d(16, 32, 5, 2, 2, bias=True)
self.conv3 = nn.Conv2d(32, 32, 5, 2, 2, bias=True)
self.fc4 = nn.Linear(s_h8 * s_h8 * 32 + nos_dim, 800, bias=True)
self.fc5 = nn.Linear(800, z_dim, bias=True)
self.nos_encode = Identity() if not enc_noise \
else MLP(input_dim=noise_dim, hidden_dim=h_dim, output_dim=h_dim, nonlinearity=nonlinearity, num_hidden_layers=2, use_nonlinearity_output=True)
def initialize_torchvision_model(config, name, d_out, **kwargs):
# get constructor and last layer names
if name == 'wideresnet50':
constructor_name = 'wide_resnet50_2'
last_layer_name = 'fc'
elif name == 'densenet121':
constructor_name = name
last_layer_name = 'classifier'
elif name in ('resnet50', 'resnet34', 'resnet18'):
constructor_name = name
last_layer_name = 'fc'
else:
raise ValueError(f'Torchvision model {name} not recognized')
# construct the default model, which has the default last layer
constructor = getattr(torchvision.models, constructor_name)
model = constructor(**kwargs)
if config.resnet_byol_path is not None:
state_dict = torch.load(config.resnet_byol_path)
model.load_state_dict(state_dict)
if config.freeze_pretrained:
for p in model.parameters():
p.requires_grad = False
freeze_batch_norm_running_params(model)
# adjust the last layer
d = getattr(model, last_layer_name).in_features
if d_out is None: # want to initialize a featurizer model
last_layer = Identity(d)
model.d_out = d
else: # want to initialize a classifier for a particular num_classes
last_layer = nn.Linear(d, d_out)
model.d_out = d_out
model.n_outputs = d
setattr(model, last_layer_name, last_layer)
# set the feature dimension as an attribute for convenience
return model
def __init__(
self,
input_dim=2, #10,
h_dim=128,
context_dim=2,
std=0.01,
num_hidden_layers=1,
nonlinearity='softplus',
noise_type='gaussian',
enc_input=False, # True,
enc_ctx=False, #True,
#init=True,
std_method='default',
):
super().__init__()
self.input_dim = input_dim
self.h_dim = h_dim
self.context_dim = context_dim
self.std = std
self.num_hidden_layers = num_hidden_layers
self.nonlinearity = nonlinearity
self.noise_type = noise_type
self.enc_input = enc_input
if self.enc_input:
inp_dim = h_dim
else:
inp_dim = input_dim
self.enc_ctx = enc_ctx
if self.enc_ctx:
ctx_dim = h_dim
else:
ctx_dim = context_dim
self.ctx_dim = ctx_dim
self.ctx_encode = Identity() if not self.enc_ctx \
else MLP(context_dim, h_dim, h_dim, nonlinearity=nonlinearity, num_hidden_layers=num_hidden_layers-1 if self.enc_ctx == 'deep' else 1, use_nonlinearity_output=True)
self.inp_encode = Identity() if not self.enc_input \
else MLP(input_dim, h_dim, h_dim, nonlinearity=nonlinearity, num_hidden_layers=num_hidden_layers-1, use_nonlinearity_output=True)
self.neglogprob = MLP(inp_dim + ctx_dim + 1,
h_dim,
1,
nonlinearity=nonlinearity,
num_hidden_layers=num_hidden_layers,
use_nonlinearity_output=False)
def __init__(
self,
input_dim=2, #10,
h_dim=128,
context_dim=2,
std=0.01,
num_hidden_layers=1,
nonlinearity='tanh',
noise_type='gaussian',
enc_input=True,
enc_ctx=True,
#init=True,
):
super().__init__()
self.input_dim = input_dim
self.h_dim = h_dim
self.context_dim = context_dim
self.std = std
self.num_hidden_layers = num_hidden_layers
self.nonlinearity = nonlinearity
self.noise_type = noise_type
self.enc_input = enc_input
if self.enc_input:
inp_dim = h_dim
else:
inp_dim = input_dim
self.enc_ctx = enc_ctx
if self.enc_ctx:
ctx_dim = h_dim
else:
ctx_dim = context_dim
#self.init = init
self.ctx_encode = Identity() if not self.enc_ctx \
else MLP(context_dim, h_dim, h_dim, nonlinearity=nonlinearity, num_hidden_layers=num_hidden_layers-1, use_nonlinearity_output=True)
self.inp_encode = Identity() if not self.enc_input \
else MLP(input_dim, h_dim, h_dim, nonlinearity=nonlinearity, num_hidden_layers=num_hidden_layers-1, use_nonlinearity_output=True)
self.dae = MLP(inp_dim + ctx_dim,
h_dim,
input_dim,
nonlinearity=nonlinearity,
num_hidden_layers=num_hidden_layers,
use_nonlinearity_output=False)
def __init__(
self,
input_dim=2,
noise_dim=2,
h_dim=64,
z_dim=2,
nonlinearity='tanh',
num_hidden_layers=1,
enc_input=False,
enc_noise=False,
clip_logvar=None,
):
super().__init__(
input_dim=input_dim,
noise_dim=noise_dim,
h_dim=h_dim,
z_dim=z_dim,
nonlinearity=nonlinearity,
num_hidden_layers=num_hidden_layers,
enc_input=enc_input,
enc_noise=enc_noise,
clip_logvar=clip_logvar,
)
inp_dim = input_dim if not enc_input else h_dim
ctx_dim = noise_dim if not enc_noise else h_dim
self.inp_encode = Identity() if not enc_input \
else MLP(input_dim=input_dim, hidden_dim=h_dim, output_dim=h_dim, nonlinearity=nonlinearity, num_hidden_layers=num_hidden_layers-1, use_nonlinearity_output=True)
self.nos_encode = Identity() if not enc_noise \
else MLP(input_dim=noise_dim, hidden_dim=h_dim, output_dim=h_dim, nonlinearity=nonlinearity, num_hidden_layers=num_hidden_layers-1, use_nonlinearity_output=True)
self.fc = MLP(input_dim=inp_dim + ctx_dim,
hidden_dim=h_dim,
output_dim=h_dim,
nonlinearity=nonlinearity,
num_hidden_layers=num_hidden_layers - 1,
use_nonlinearity_output=True)
self.reparam = NormalDistributionLinear(h_dim,
z_dim,
nonlinearity=clip_logvar)
def __init__(
self,
input_dim=2,
z_dim=2,
noise_dim=2,
h_dim=64,
nonlinearity='tanh',
num_hidden_layers=1,
enc_input=False,
enc_latent=False,
clip_logvar=None,
):
super().__init__()
self.input_dim = input_dim
self.z_dim = z_dim
self.noise_dim = noise_dim
self.h_dim = h_dim
self.nonlinearity = nonlinearity
self.num_hidden_layers = num_hidden_layers
self.enc_input = enc_input
self.enc_latent = enc_latent
inp_dim = input_dim if not enc_input else h_dim
ltt_dim = z_dim if not enc_latent else h_dim
self.inp_encode = Identity() if not enc_input \
else MLP(input_dim=input_dim, hidden_dim=h_dim, output_dim=h_dim, nonlinearity=nonlinearity, num_hidden_layers=num_hidden_layers-1, use_nonlinearity_output=True)
self.ltt_encode = Identity() if not enc_latent \
else MLP(input_dim=z_dim, hidden_dim=h_dim, output_dim=h_dim, nonlinearity=nonlinearity, num_hidden_layers=num_hidden_layers-1, use_nonlinearity_output=True)
self.fc = MLP(input_dim=inp_dim + ltt_dim,
hidden_dim=h_dim,
output_dim=h_dim,
nonlinearity=nonlinearity,
num_hidden_layers=num_hidden_layers - 1,
use_nonlinearity_output=True)
self.reparam = NormalDistributionLinear(h_dim,
noise_dim,
nonlinearity=clip_logvar)
def __init__(
self,
input_dim=2,
noise_dim=2,
h_dim=64,
z_dim=2,
nonlinearity='softplus',
num_hidden_layers=1,
std=1.,
init='none', #'gaussian',
enc_noise=False,
):
super().__init__(
input_dim=input_dim,
noise_dim=noise_dim,
h_dim=h_dim,
z_dim=z_dim,
nonlinearity=nonlinearity,
num_hidden_layers=num_hidden_layers,
std=std,
init=init,
enc_noise=enc_noise,
)
nos_dim = noise_dim if not enc_noise else h_dim
self.inp_encode = MLP(input_dim=input_dim,
hidden_dim=h_dim,
output_dim=h_dim,
nonlinearity=nonlinearity,
num_hidden_layers=num_hidden_layers,
use_nonlinearity_output=True)
self.nos_encode = Identity() if not enc_noise \
else MLP(input_dim=noise_dim, hidden_dim=h_dim, output_dim=h_dim, nonlinearity=nonlinearity, num_hidden_layers=0, use_nonlinearity_output=True)
self.fc = MLP(input_dim=h_dim + nos_dim,
hidden_dim=h_dim,
output_dim=z_dim,
nonlinearity=nonlinearity,
num_hidden_layers=1,
use_nonlinearity_output=False)
if self.init == 'gaussian':
self.reset_parameters()
else:
pass
class ConditionalARDAE(nn.Module):
def __init__(
self,
input_dim=2, #10,
h_dim=128,
context_dim=2,
std=0.01,
num_hidden_layers=1,
nonlinearity='tanh',
noise_type='gaussian',
enc_input=True,
enc_ctx=True,
#init=True,
std_method='default',
):
super().__init__()
self.input_dim = input_dim
self.h_dim = h_dim
self.context_dim = context_dim
self.std = std
self.num_hidden_layers = num_hidden_layers
self.nonlinearity = nonlinearity
self.noise_type = noise_type
self.enc_input = enc_input
if self.enc_input:
inp_dim = h_dim
else:
inp_dim = input_dim
self.enc_ctx = enc_ctx
if self.enc_ctx:
ctx_dim = h_dim
else:
ctx_dim = context_dim
self.ctx_encode = Identity() if not self.enc_ctx \
else MLP(context_dim, h_dim, h_dim, nonlinearity=nonlinearity, num_hidden_layers=num_hidden_layers-1, use_nonlinearity_output=True)
self.inp_encode = Identity() if not self.enc_input \
else MLP(input_dim, h_dim, h_dim, nonlinearity=nonlinearity, num_hidden_layers=num_hidden_layers-1, use_nonlinearity_output=True)
self.neglogprob = MLP(inp_dim + ctx_dim + 1,
h_dim,
1,
nonlinearity=nonlinearity,
num_hidden_layers=num_hidden_layers,
use_nonlinearity_output=False)
def reset_parameters(self):
nn.init.normal_(self.neglogprob.fc.weight)
self.inp_encode.apply(weights_init)
def add_noise(self, input, std=None):
std = self.std if std is None else std
if self.noise_type == 'gaussian':
return add_gaussian_noise(input, std)
elif self.noise_type == 'uniform':
return add_uniform_noise(input, std)
elif self.noise_type == 'laplace':
return add_laplace_noise(input, std)
else:
raise NotImplementedError
def loss(self, input, target):
# recon loss (likelihood)
recon_loss = F.mse_loss(input, target) #, reduction='sum')
return recon_loss
def forward(self, input, context, std=None, scale=None):
# init
assert input.dim() == 3 # bsz x ssz x x_dim
assert context.dim() == 3 # bsz x 1 x ctx_dim
batch_size = input.size(0)
sample_size = input.size(1)
if std is None:
std = input.new_zeros(batch_size, sample_size, 1)
else:
assert torch.is_tensor(std)
if scale is None:
scale = 1.
# reschape
input = input.view(batch_size * sample_size,
self.input_dim) # bsz*ssz x xdim
_, context = expand_tensor(context,
sample_size=sample_size,
do_unsqueeze=False) # bsz*ssz x xdim
#context = context.view(batch_size*sample_size, -1) # bsz*ssz x xdim
std = std.view(batch_size * sample_size, 1)
# add noise
x_bar, eps = self.add_noise(input, std)
# grad true
x_bar.requires_grad = True
# encode
ctx = self.ctx_encode(context)
inp = self.inp_encode(x_bar)
# concat
h = torch.cat([inp, ctx, std], dim=1)
# (unnorm) logprob
logprob = -self.neglogprob(h)
logprob = torch.sum(logprob)
# de-noise with context
glogprob = grad(logprob, x_bar)
''' get loss '''
#loss = (std**2)*self.loss(std*glogprob, -eps)
loss = self.loss(std * glogprob, -eps)
# return
return None, loss
def glogprob(self, input, context, std=None, scale=None):
# init
assert input.dim() == 3 # bsz x ssz x x_dim
assert context.dim() == 3 # bsz x 1 x ctx_dim
#std = self.std if std is None else std
batch_size = input.size(0)
sample_size = input.size(1)
if std is None:
std = input.new_zeros(batch_size * sample_size, 1)
else:
assert torch.is_tensor(std)
if scale is None:
scale = 1.
# reschape
input = input.view(batch_size * sample_size,
self.input_dim) # bsz*ssz x xdim
_, context = expand_tensor(context,
sample_size=sample_size,
do_unsqueeze=False) # bsz*ssz x xdim
#context = context.view(batch_size*sample_size, -1) # bsz*ssz x xdim
std = std.view(batch_size * sample_size, 1)
# grad true
input.requires_grad = True
# encode
ctx = self.ctx_encode(context)
inp = self.inp_encode(input)
# concat
h = torch.cat([inp, ctx, std], dim=1)
# (unnorm) logprob
logprob = -self.neglogprob(h)
logprob = torch.sum(logprob)
# de-noise with context
glogprob = grad(logprob, input)
return glogprob.view(batch_size, sample_size, self.input_dim)
def __init__(self,
num_inputs,
num_actions,
hidden_dim,
noise_dim,
num_enc_layers,
num_fc_layers,
args,
nonlinearity='elu',
fc_type='mlp'):
super(StochasticPolicy, self).__init__()
self.num_enc_layers = num_enc_layers
self.num_fc_layers = num_fc_layers
self.args = args
self.num_actions = num_actions
self.noise_dim = noise_dim
assert noise_dim >= num_actions, 'noise_dim: {}, num_actions: {}'.format(
noise_dim, num_actions)
inp_dim = num_inputs if num_enc_layers < 0 else hidden_dim
if num_enc_layers < 0:
self.encode = Identity()
else:
self.encode = MLP(num_inputs,
hidden_dim,
hidden_dim,
nonlinearity=nonlinearity,
num_hidden_layers=num_enc_layers,
use_nonlinearity_output=True)
if fc_type == 'mlp':
self.fc = MLP(inp_dim + noise_dim,
hidden_dim,
num_actions,
nonlinearity=nonlinearity,
num_hidden_layers=num_fc_layers,
use_nonlinearity_output=False)
torch.nn.init.normal_(self.fc.fc.weight, std=1.)
elif fc_type == 'wnres':
self.fc = ResMLP(inp_dim + noise_dim,
hidden_dim,
num_actions,
nonlinearity=nonlinearity,
num_hidden_layers=num_fc_layers,
layer='wnlinear',
use_nonlinearity_output=False)
elif fc_type == 'res':
self.fc = ResMLP(inp_dim + noise_dim,
hidden_dim,
num_actions,
nonlinearity=nonlinearity,
num_hidden_layers=num_fc_layers,
layer='linear',
use_nonlinearity_output=False)
elif fc_type == 'mlpdeep':
self.fc = MLP(inp_dim + noise_dim,
64,
num_actions,
nonlinearity=nonlinearity,
num_hidden_layers=num_fc_layers,
use_nonlinearity_output=False)
torch.nn.init.normal_(self.fc.fc.weight, std=1.)
elif fc_type == 'wnresdeep':
self.fc = ResMLP(inp_dim + noise_dim,
64,
num_actions,
nonlinearity=nonlinearity,
num_hidden_layers=num_fc_layers,
layer='wnlinear',
use_nonlinearity_output=False)
elif fc_type == 'resdeep':
self.fc = ResMLP(inp_dim + noise_dim,
64,
num_actions,
nonlinearity=nonlinearity,
num_hidden_layers=num_fc_layers,
layer='linear',
use_nonlinearity_output=False)
else:
raise NotImplementedError
def __init__(
self,
noise_dim=100,
z_dim=32,
c_dim=512, #450,
h_dim=800,
num_hidden_layers=1,
nonlinearity='elu', #act=nn.ELU(),
do_center=False,
enc_noise=False,
enc_type='mlp',
):
super().__init__()
self.noise_dim = noise_dim
self.z_dim = z_dim
self.c_dim = c_dim
self.h_dim = h_dim
self.num_hidden_layers = num_hidden_layers
assert num_hidden_layers > 0
self.nonlinearity = nonlinearity
self.do_center = do_center
self.enc_noise = enc_noise
nos_dim = noise_dim if not enc_noise else c_dim
self.enc_type = enc_type
assert enc_type in [
'mlp', 'res-wn-mlp', 'res-mlp', 'res-wn-mlp-lin', 'res-mlp-lin'
]
act = get_afunc(nonlinearity_type=nonlinearity)
self.inp_encode = nn.Sequential(
nn_.ResConv2d(1, 16, 3, 2, padding=1, activation=act), act,
nn_.ResConv2d(16, 16, 3, 1, padding=1, activation=act), act,
nn_.ResConv2d(16, 32, 3, 2, padding=1, activation=act), act,
nn_.ResConv2d(32, 32, 3, 1, padding=1, activation=act), act,
nn_.ResConv2d(32, 32, 3, 2, padding=1, activation=act), act,
nn_.Reshape((-1, 32 * 4 * 4)), nn_.ResLinear(32 * 4 * 4, c_dim),
act)
#self.fc = nn.Sequential(
# nn.Linear(c_dim + nos_dim, h_dim, bias=True),
# act,
# nn.Linear(h_dim, z_dim, bias=True),
# )
if enc_type == 'mlp':
self.fc = MLP(input_dim=c_dim + nos_dim,
hidden_dim=h_dim,
output_dim=z_dim,
nonlinearity=nonlinearity,
num_hidden_layers=num_hidden_layers,
use_nonlinearity_output=False)
elif enc_type == 'res-wn-mlp':
self.fc = ResMLP(input_dim=c_dim + nos_dim,
hidden_dim=h_dim,
output_dim=z_dim,
nonlinearity=nonlinearity,
num_hidden_layers=num_hidden_layers,
use_nonlinearity_output=False,
layer='wnlinear')
elif enc_type == 'res-mlp':
self.fc = ResMLP(input_dim=c_dim + nos_dim,
hidden_dim=h_dim,
output_dim=z_dim,
nonlinearity=nonlinearity,
num_hidden_layers=num_hidden_layers,
use_nonlinearity_output=False,
layer='linear')
elif enc_type == 'res-wn-mlp-lin':
self.fc = nn.Sequential(
ResMLP(input_dim=c_dim + nos_dim,
hidden_dim=h_dim,
output_dim=h_dim,
nonlinearity=nonlinearity,
num_hidden_layers=num_hidden_layers - 1,
use_nonlinearity_output=True,
layer='wnlinear'),
nn.Linear(h_dim, z_dim, bias=True),
)
elif enc_type == 'res-mlp-lin':
self.fc = nn.Sequential(
ResMLP(input_dim=c_dim + nos_dim,
hidden_dim=h_dim,
output_dim=h_dim,
nonlinearity=nonlinearity,
num_hidden_layers=num_hidden_layers - 1,
use_nonlinearity_output=True,
layer='linear'),
nn.Linear(h_dim, z_dim, bias=True),
)
else:
raise NotImplementedError
self.nos_encode = Identity() if not enc_noise \
else nn.Sequential(
nn.Linear(noise_dim, c_dim, bias=True),
act,
)