def __init__(self, dataset, encoder, decoder, volsampler, colorcal, dt, stepjitter=0.01, estimatebg=False):
super(Autoencoder, self).__init__()
self.estimatebg = estimatebg
self.allcameras = dataset.get_allcameras()
self.encoder = encoder
self.decoder = decoder
self.volsampler = volsampler
self.bg = nn.ParameterDict({
k: nn.Parameter(torch.ones(3, v["size"][1], v["size"][0]), requires_grad=estimatebg)
for k, v in dataset.get_krt().items()})
self.colorcal = colorcal
self.dt = dt
self.stepjitter = stepjitter
self.imagemean = dataset.imagemean
self.imagestd = dataset.imagestd
if dataset.known_background():
dataset.get_background(self.bg)
def __init__(self, operation: MixedOperation, memo: dict[str, Any],
mutate_kwargs: dict[str, Any]) -> None:
# Sampling arguments. This should have the same keys with `operation.mutable_arguments`
operation._arch_alpha = nn.ParameterDict()
for name, spec in operation.search_space_spec().items():
if name in memo:
alpha = memo[name]
if len(alpha) != spec.size:
raise ValueError(
f'Architecture parameter size of same label {name} conflict: {len(alpha)} vs. {spec.size}'
)
else:
alpha = nn.Parameter(torch.randn(spec.size) * 1E-3)
operation._arch_alpha[name] = alpha
operation.parameters = functools.partial(self.parameters,
module=operation) # bind self
operation.named_parameters = functools.partial(self.named_parameters,
module=operation)
operation._softmax = mutate_kwargs.get('softmax', nn.Softmax(-1))
def __init__(self, input_dim, hidden_dim):
# initialize all pytorch class
super(NADE, self).__init__()
# store the dimension input and output
self.D = input_dim
self.H = hidden_dim
# store the parameters of the network
self.params = nn.ParameterDict({
"V":
nn.Parameter(torch.randn(self.D, self.H)),
"b":
nn.Parameter(torch.zeros(self.D)),
"W":
nn.Parameter(torch.randn(self.H, self.D)),
"c":
nn.Parameter(torch.zeros(1, self.H)),
})
# initialize the parameters
nn.init.xavier_normal_(self.params["V"])
nn.init.xavier_normal_(self.params["W"])
def __init__(self,
input_size=64,
num_blocks=3,
growth_rate=14,
num_fin_growth=1,
feature_mode=None,
drop_rate=0.2,
bare=False):
super(ConvModule, self).__init__()
self.num_blocks = num_blocks
self.drop_rate = drop_rate
if not bare:
self.num_blocks = num_blocks
self.params = nn.ParameterDict()
num_channels = [
max(1, nci * growth_rate) for nci in range(num_blocks)
]
#num_channels.insert(0, num_init_growth*growth_rate)
num_channels.append(num_fin_growth * growth_rate)
feature_size = input_size
for nci in range(num_blocks):
prefix = 'l' + str(nci)
conv_size = 3
#if feature_size>5: feature_size /= 2
#else: feature_size -= conv_size - 1
feature_size /= 2
if nci == num_blocks - 1:
if feature_mode == 'final_fuse':
conv_size = int(input_size / 2**nci)
feature_size = 1
elif feature_mode == 'avg_pool':
feature_size = 1
stdv = 1. / math.sqrt(num_channels[nci] * conv_size**3)
self.params[prefix + '_conv_w'] = nn.Parameter(
torch.Tensor(num_channels[nci + 1], num_channels[nci],
conv_size, conv_size,
conv_size).uniform_(-stdv, stdv))
self.params[prefix + '_conv_b'] = nn.Parameter(
torch.Tensor(num_channels[nci + 1]).uniform_(-stdv, stdv))
self.end_feature_size = feature_size
def __init__(self, num_obj=1, with_rotvec=False):
super(StatePredictorParams, self).__init__()
maxT = 100
### initialize parameters for hand
# 0:3 hand end-effector position is at x,y,z
# 3 hand azimuth
# 4 hand elevation
init_params = 0.1 * torch.randn(1, maxT, 5) # BS x T x n_param
# hand azimuth angle initial values can center at pi/2 for easier transition to 0 or pi
# angle around 0 used for push L --> R / pull R --> L
# angle around pi used for push R --> L / pull L --> R
init_params[0, :, 3] += np.pi / 2
hand = nn.Parameter(init_params)
### initialize parameters for objects (always on table)
# 0:3 object0 size in x,y,z dimensions
# 3:6 object0 position is at x,y,z (z is on the table for obj0)
# 6:9 object0 rotation vector
init_params = 0.1 * torch.randn(1, maxT, 6) # BS x T x n_param
if with_rotvec:
init_params = torch.cat(
(init_params, 0.01 * torch.randn(1, maxT, 3)), dim=2)
obj0 = nn.Parameter(init_params)
param_dict = {'hand': hand, 'obj0': obj0}
### initialize parameters for obj1
if num_obj == 2:
# 0:3 object1 size in x,y,z dimensions
# 3:6 object1 position is at x,y,z (z is on the table for obj0)
# 6:9 object1 rotation vector
init_params = 0.1 * torch.randn(1, maxT, 9) # BS x T x n_param
if with_rotvec:
init_params = torch.cat(
(init_params, 0.01 * torch.randn(1, maxT, 3)), dim=2)
obj1 = nn.Parameter(init_params)
param_dict['obj1'] = obj1
# make states into parameter dictionary
self.states = nn.ParameterDict(param_dict)
def __init__(self, useCuda=True, num_groups=3, num_classes=None):
super(WideResNet_50_2, self).__init__()
self.num_groups = num_groups
self.num_classes = num_classes
params = model_zoo.load_url(
'https://s3.amazonaws.com/modelzoo-networks/wide-resnet-50-2-export-5ae25d50.pth'
)
# convert numpy arrays to torch Variables
print('loaded weigths shape...')
for k, v in sorted(params.items()):
print(k, tuple(v.shape))
if useCuda:
params[k] = Variable(v.cuda(), requires_grad=True)
else:
params[k] = Variable(v, requires_grad=True)
# Change the last fully connected layer
if num_classes is not None:
tmp = nn.Linear([256, 512, 1024, 2048][num_groups], num_classes)
if useCuda:
tmp = tmp.cuda()
params['fc_weight'] = tmp.weight
params['fc_bias'] = tmp.bias
print('\nTotal parameters:', sum(v.numel() for v in params.values()))
# replace all '.' in the dictionary by '_'
params_tmp = {}
for key in params.keys():
params_tmp[key.replace('.', '_')] = params[key]
params = params_tmp
del params_tmp
#create nn.Parameters to return them with parameters() function.
self.params = nn.ParameterDict({})
for key in params.keys():
self.params.update(
{key: nn.Parameter(params[key], requires_grad=True)})
def __init__(self, categories):
super().__init__()
logging.info(f'creating model with categories: {categories}')
# todo(will.brennan) - find a nicer way of saving the categories in the state dict...
self._categories = nn.ParameterDict({i: nn.Parameter(torch.Tensor(0)) for i in categories})
num_categories = len(self._categories)
self.semantic = SemanticBranch()
self.detail = DetailBranch()
self.bga = BilateralGuidedAggregationLayer(128)
seg_head_names = ['out', 'aux_c2', 'aux_c3', 'aux_c4', 'aux_c5']
self.seg_heads = nn.ModuleDict(
{
'out': SegmentationHead(128, num_categories, scale_factor=8),
'aux_c2': SegmentationHead(16, num_categories, scale_factor=4),
'aux_c3': SegmentationHead(32, num_categories, scale_factor=8),
'aux_c4': SegmentationHead(64, num_categories, scale_factor=16),
'aux_c5': SegmentationHead(128, num_categories, scale_factor=32),
}
)
def __init__(self, seq_length, input_dim, hidden_dim, num_classes,
batch_size, device):
super(peepLSTM, self).__init__()
########################
# PUT YOUR CODE HERE #
#######################
self.seq_length = seq_length
self.input_dim = input_dim
self.hidden_dim = hidden_dim
self.num_classes = num_classes
self.batch_size = batch_size
self.device = device
self.params = nn.ParameterDict()
for gate in ['i', 'f', 'o']:
self.params['W_' + gate + 'h'] = nn.Parameter(
torch.empty(hidden_dim, hidden_dim))
self.params['W_' + gate + 'x'] = nn.Parameter(
torch.empty(input_dim, hidden_dim))
nn.init.kaiming_normal_(self.params['W_' + gate + 'x'],
nonlinearity='linear')
nn.init.kaiming_normal_(self.params['W_' + gate + 'h'],
nonlinearity='linear')
self.params['b_' + gate] = nn.Parameter(torch.zeros(1, hidden_dim))
self.params['W_cx'] = nn.Parameter(torch.empty(input_dim, hidden_dim))
self.params['b_c'] = nn.Parameter(torch.zeros(1, hidden_dim))
nn.init.kaiming_normal_(self.params['W_cx'], nonlinearity='linear')
self.params['W_ph'] = nn.Parameter(torch.empty(hidden_dim,
num_classes))
nn.init.kaiming_normal_(self.params['W_ph'], nonlinearity='linear')
self.params['b_p'] = nn.Parameter(torch.zeros(1, num_classes))
self.tanh = nn.Tanh()
self.sigmoid = nn.Sigmoid()
self.softmax = nn.LogSoftmax(dim=-1)
self.embedding = nn.Embedding(3, input_dim)
def __init__(self, model, lstm_size=64, lstm_num_layers=1, tanh_constant=1.5, cell_exit_extra_step=False,
skip_target=0.4, temperature=None, branch_bias=0.25, entropy_reduction="sum"):
super().__init__(model)
self.lstm_size = lstm_size
self.lstm_num_layers = lstm_num_layers
self.tanh_constant = tanh_constant
self.temperature = temperature
self.cell_exit_extra_step = cell_exit_extra_step
self.skip_target = skip_target
self.branch_bias = branch_bias
self.lstm = StackedLSTMCell(self.lstm_num_layers, self.lstm_size, False)
self.attn_anchor = nn.Linear(self.lstm_size, self.lstm_size, bias=False)
self.attn_query = nn.Linear(self.lstm_size, self.lstm_size, bias=False)
self.v_attn = nn.Linear(self.lstm_size, 1, bias=False)
self.g_emb = nn.Parameter(torch.randn(1, self.lstm_size) * 0.1)
self.skip_targets = nn.Parameter(torch.tensor([1.0 - self.skip_target, self.skip_target]), requires_grad=False) # pylint: disable=not-callable
assert entropy_reduction in ["sum", "mean"], "Entropy reduction must be one of sum and mean."
self.entropy_reduction = torch.sum if entropy_reduction == "sum" else torch.mean
self.cross_entropy_loss = nn.CrossEntropyLoss(reduction="none")
self.bias_dict = nn.ParameterDict()
self.max_layer_choice = 0
for mutable in self.mutables:
if isinstance(mutable, LayerChoice):
if self.max_layer_choice == 0:
self.max_layer_choice = len(mutable)
assert self.max_layer_choice == len(mutable), \
"ENAS mutator requires all layer choice have the same number of candidates."
# We are judging by keys and module types to add biases to layer choices. Needs refactor.
if "reduce" in mutable.key:
def is_conv(choice):
return "conv" in str(type(choice)).lower()
bias = torch.tensor([self.branch_bias if is_conv(choice) else -self.branch_bias # pylint: disable=not-callable
for choice in mutable])
self.bias_dict[mutable.key] = nn.Parameter(bias, requires_grad=False)
self.embedding = nn.Embedding(self.max_layer_choice + 1, self.lstm_size)
self.soft = nn.Linear(self.lstm_size, self.max_layer_choice, bias=False)
def __init__(self, emb_dim, domain_dims):
super(APE, self).__init__()
self.save_path = None
self.num_domains = len(domain_dims)
self.emb_dim = emb_dim
self.num_entities = sum(domain_dims.values())
self.emb_layer = nn.Embedding(num_embeddings=self.num_entities,
embedding_dim=emb_dim)
self.c = nn.Parameter(torch.from_numpy(np.random.random(1)))
self.mode = 'train'
k = 0
self.pair_W = nn.ParameterDict({})
for i in range(self.num_domains):
for j in range(i + 1, self.num_domains):
w_k = nn.Parameter(data=FT(np.random.random(1)))
self.pair_W[str(k)] = w_k
k += 1
return
def __init__(self, n, loss_fn, lr=1e-3, **modelParams):
super(MetaSGD, self).__init__()
self.DataParallel = modelParams[
'data_parallel'] if 'data_parallel' in modelParams else False
#######################################################
# For CNN only
modelParams['dims'] = [modelParams['embed_size'], 64, 128, 256, 256]
modelParams['kernel_sizes'] = [3, 3, 3, 3]
modelParams['paddings'] = [1, 1, 1, 1]
modelParams['relus'] = [True, True, True, True]
modelParams['pools'] = ['max', 'max', 'max', 'ada']
#######################################################
self.Learner = BaseLearner(n, **modelParams) # 基学习器内部含有beta
self.Alpha = nn.ParameterDict({
rename(name):
nn.Parameter(lr * t.ones_like(val, requires_grad=True))
for name, val in self.Learner.named_parameters()
}) # 初始化alpha
self.LossFn = loss_fn
def __init__(self, ntype_dict, etypes, in_size, hidden_size, out_size,
n_layers, embedding_size):
super(HeteroRGCN, self).__init__()
# Use trainable node embeddings as featureless inputs.
embed_dict = {
ntype: nn.Parameter(torch.Tensor(num_nodes, in_size))
for ntype, num_nodes in ntype_dict.items() if ntype != 'target'
}
for key, embed in embed_dict.items():
nn.init.xavier_uniform_(embed)
self.embed = nn.ParameterDict(embed_dict)
# create layers
self.layers = nn.ModuleList()
self.layers.append(HeteroRGCNLayer(embedding_size, hidden_size,
etypes))
# hidden layers
for i in range(n_layers - 1):
self.layers.append(
HeteroRGCNLayer(hidden_size, hidden_size, etypes))
# output layer
self.layers.append(nn.Linear(hidden_size, out_size))
def __init__(self,
rating_vals,
user_in_units,
movie_in_units,
msg_units,
out_units,
dropout_rate=0.0,
agg='stack', # or 'sum'
agg_act=None,
out_act=None,
share_user_item_param=False):
super(GCMCLayer, self).__init__()
self.rating_vals = rating_vals
self.agg = agg
self.share_user_item_param = share_user_item_param
self.ufc = nn.Linear(msg_units, out_units)
if share_user_item_param:
self.ifc = self.ufc
else:
self.ifc = nn.Linear(msg_units, out_units)
if agg == 'stack':
# divide the original msg unit size by number of ratings to keep
# the dimensionality
assert msg_units % len(rating_vals) == 0
msg_units = msg_units // len(rating_vals)
self.dropout = nn.Dropout(dropout_rate)
self.W_r = nn.ParameterDict()
for rating in rating_vals:
# PyTorch parameter name can't contain "."
rating = str(rating).replace('.', '_')
if share_user_item_param and user_in_units == movie_in_units:
self.W_r[rating] = nn.Parameter(th.randn(user_in_units, msg_units))
self.W_r['rev-%s' % rating] = self.W_r[rating]
else:
self.W_r[rating] = nn.Parameter(th.randn(user_in_units, msg_units))
self.W_r['rev-%s' % rating] = nn.Parameter(th.randn(movie_in_units, msg_units))
self.agg_act = get_activation(agg_act)
self.out_act = get_activation(out_act)
self.reset_parameters()
def __init__(self,
num_input_features,
growth_rate,
drop_rate,
final=False,
inidev=0.01,
use_gpu=True):
super(doubconv_pool_pwise, self).__init__()
if use_gpu:
device = torch.device('cuda')
else:
device = torch.device('cpu')
self.params = nn.ParameterDict()
self.params['convw1'] = nn.Parameter(
V(torch.randn(
(growth_rate, num_input_features, 3, 3, 3), device=device) *
inidev,
requires_grad=True))
self.params['convw2'] = nn.Parameter(
V(torch.randn(
(growth_rate, growth_rate, 3, 3, 3), device=device) * inidev,
requires_grad=True))
if final:
self.params['convb2'] = nn.Parameter(
V(torch.randn((growth_rate), device=device) * inidev,
requires_grad=True))
self.params['bnw1'] = nn.Parameter(
V(torch.randn(num_input_features, device=device) * inidev,
requires_grad=True))
self.params['bnb1'] = nn.Parameter(
V(torch.randn(num_input_features, device=device) * inidev,
requires_grad=True))
self.params['bnw2'] = nn.Parameter(
V(torch.randn(growth_rate, device=device) * inidev,
requires_grad=True))
self.params['bnb1'] = nn.Parameter(
V(torch.randn(growth_rate, device=device) * inidev,
requires_grad=True))
self.drop_rate = drop_rate
def __init__(self, dim_z, Nsliders):
super(walk_embed, self).__init__()
self.dim_z = dim_z
self.Nsliders = Nsliders
self.w_embed = nn.ParameterDict({
'zoom':
nn.Parameter(
torch.Tensor(
np.random.normal(0.0, 0.02,
[15, 1, self.dim_z, Nsliders])).cuda()),
'shiftx':
nn.Parameter(
torch.Tensor(
np.random.normal(0.0, 0.02,
[15, 1, self.dim_z, Nsliders])).cuda()),
'rotate2d':
nn.Parameter(
torch.Tensor(
np.random.normal(0.0, 0.02,
[15, 1, self.dim_z, Nsliders])).cuda())
})
def __init__(self, dim_z, Nsliders):
super(walk_embed, self).__init__()
self.dim_z = dim_z
self.Nsliders = Nsliders
self.w_embed = nn.ParameterDict({
'blondhair':
nn.Parameter(
torch.Tensor(
np.random.normal(0.0, 0.02,
[6, 1, self.dim_z, Nsliders])).cuda()),
'paleskin':
nn.Parameter(
torch.Tensor(
np.random.normal(0.0, 0.02,
[6, 1, self.dim_z, Nsliders])).cuda()),
'male':
nn.Parameter(
torch.Tensor(
np.random.normal(0.0, 0.02,
[6, 1, self.dim_z, Nsliders])).cuda())
})
def __init__(self,
embed_size,
g,
bias=True,
activation=None,
self_loop=False,
dropout=0.0):
super(RelGraphConvHeteroEmbed, self).__init__()
self.embed_size = embed_size
self.g = g
self.bias = bias
self.activation = activation
self.self_loop = self_loop
# create weight embeddings for each node for each relation
self.embeds = nn.ParameterDict()
for srctype, etype, dsttype in g.canonical_etypes:
print("{} {} {}".format(srctype, etype, dsttype))
embed = nn.Parameter(
th.Tensor(g.number_of_nodes(srctype), self.embed_size))
nn.init.xavier_uniform_(embed, gain=nn.init.calculate_gain('relu'))
self.embeds["{}-{}-{}".format(srctype, etype, dsttype)] = embed
# bias
if self.bias:
self.h_bias = nn.Parameter(th.Tensor(embed_size))
nn.init.zeros_(self.h_bias)
# weight for self loop
if self.self_loop:
self.self_embeds = nn.ParameterList()
for ntype in g.ntypes:
embed = nn.Parameter(
th.Tensor(g.number_of_nodes(ntype), embed_size))
nn.init.xavier_uniform_(embed,
gain=nn.init.calculate_gain('relu'))
self.self_embeds.append(embed)
self.dropout = nn.Dropout(dropout)
def test_flatten_module(self):
a = nn.Module()
b = nn.Module()
c = nn.Module()
d = nn.Module()
pa = nn.Parameter()
pb = nn.Parameter()
pc = nn.Parameter()
pd = nn.Parameter()
nest = nn.ModuleDict({
'a': a,
'b': b,
'list': nn.ModuleList([c, d]),
'plist': nn.ParameterList([pa, pb]),
'pdict': nn.ParameterDict({
'pc': pc,
'pd': pd
})
})
flattend = alf.algorithms.algorithm._flatten_module(nest)
self.assertEqual(set(map(id, flattend)),
set(map(id, [a, b, c, d, pa, pb, pc, pd])))
def __init__(self,
G,
in_size,
hidden_size,
out_size,
use_dr_pre,
pre_v_dict,
if_non_active=False):
super(HeteroGCNNet, self).__init__()
# Use trainable node embeddings as featureless inputs.
if use_dr_pre:
embed_dict = {}
for ntype in G.ntypes:
if ntype == 'description' or ntype == 'review':
pre_vec = torch.tensor(pre_v_dict[ntype],
dtype=torch.float)
assert pre_vec.shape[0] == G.number_of_nodes(
ntype), 'The shape of pre_vec is wrong!'
embed = nn.Parameter(pre_vec)
embed_dict[ntype] = embed
else:
embed = nn.Parameter(
torch.zeros(G.number_of_nodes(ntype), in_size))
nn.init.xavier_uniform_(embed)
embed_dict[ntype] = embed
else:
embed_dict = {}
for ntype in G.ntypes:
embed = nn.Parameter(
torch.zeros(G.number_of_nodes(ntype), in_size))
nn.init.xavier_uniform_(embed)
embed_dict[ntype] = embed
self.embed = nn.ParameterDict(embed_dict)
# create layers
self.layer1 = HeteroGCNLayer(in_size, hidden_size, G.etypes)
self.layer2 = HeteroGCNLayer(hidden_size, out_size, G.etypes)
self.if_non_active = if_non_active
self.h_dict = None
def __init__(self):
super(Test_entNet_Model, self).__init__()
self.device = "cuda"
# embedding parameters
self.H = nn.init.normal_(torch.empty(H_size, requires_grad=True, dtype=torch.float, device=self.device), mean=0.0, std=0.1)
self.W = nn.init.normal_(torch.empty(W_size, requires_grad=True, dtype=torch.float, device=self.device), mean=0.0, std=0.1)
self.params = nn.ParameterDict({
'F': nn.Parameter(nn.init.normal_(torch.empty(input_size, requires_grad=True, dtype=torch.float, device=self.device), mean=0.0, std=0.1)),
# shared parameters
'X': nn.Parameter(nn.init.normal_(torch.empty(X_size, requires_grad=True, dtype=torch.float, device=self.device), mean=0.0, std=0.1)),
'Y': nn.Parameter(nn.init.normal_(torch.empty(Y_size, requires_grad=True, dtype=torch.float, device=self.device), mean=0.0, std=0.1)),
'Z': nn.Parameter(nn.init.normal_(torch.empty(Z_size, requires_grad=True, dtype=torch.float, device=self.device), mean=0.0, std=0.1)),
# answer parameters
'R': nn.Parameter(nn.init.normal_(torch.empty(R_size, requires_grad=True, dtype=torch.float, device=self.device), mean=0.0, std=0.1)),
'K': nn.Parameter(nn.init.normal_(torch.empty(K_size, requires_grad=True, dtype=torch.float, device=self.device), mean=0.0, std=0.1))
})
# dropout
self.dropout = nn.Dropout(p=0.3)
def __init__(self, input_dim, hidden_dim):
"""Initializes a new NADE instance.
Args:
input_dim: The dimension of the input.
hidden_dim: The dimmension of the hidden layer. NADE only supports one
hidden layer.
"""
super().__init__()
self._input_dim = input_dim
self._hidden_dim = hidden_dim
self.params = nn.ParameterDict({
'in_W':
nn.Parameter(torch.zeros(self._hidden_dim, self._input_dim)),
'in_b':
nn.Parameter(torch.zeros(self._hidden_dim, )),
'h_W':
nn.Parameter(torch.zeros(self._input_dim, self._hidden_dim)),
'h_b':
nn.Parameter(torch.zeros(self._input_dim, )),
})
nn.init.kaiming_normal_(self.params['in_W'])
nn.init.kaiming_normal_(self.params['h_W'])
def __init__(self, state_dim, action_dim):
super(CriticNet, self).__init__()
self.state_dim = state_dim
self.action_dim = action_dim
num_hiddens = 300
self.params = nn.ParameterDict({
'w1_s':
nn.Parameter(torch.randn(self.state_dim, num_hiddens) * 0.001),
'w1_a':
nn.Parameter(torch.randn(action_dim, num_hiddens) * 0.001),
'b1':
nn.Parameter(torch.zeros(1, num_hiddens))
})
self.linear = nn.Linear(num_hiddens, 1)
for name, params in self.linear.named_parameters():
if 'bias' in name:
init.constant_(params, val=0.)
else:
init.normal_(params, mean=0., std=0.001)
def __init__(self,
fiber,
nonlin=nn.ReLU(inplace=True),
num_layers: int = 0):
"""Initializer.
Args:
fiber: Fiber() of feature multiplicities and types
nonlin: nonlinearity to use everywhere
num_layers: non-negative number of linear layers in fnc
"""
super().__init__()
self.fiber = fiber
self.nonlin = nonlin
self.num_layers = num_layers
# Regularization for computing phase: gradients explode otherwise
self.eps = 1e-12
# Norm mappings: 1 per feature type
self.bias = nn.ParameterDict()
for m, d in self.fiber.structure:
self.bias[str(d)] = nn.Parameter(torch.randn(m).view(1, m))
def __init__(self, categories):
super().__init__()
logging.info(f'creating model with categories: {categories}')
# todo(will.brennan) - find a nicer way of saving the categories in the state dict...
self._categories = nn.ParameterDict(
{i: nn.Parameter(torch.Tensor(0))
for i in categories})
num_categories = len(self._categories)
self.model = detection.maskrcnn_resnet50_fpn(pretrained=True)
logging.debug('changing num_categories for bbox predictor')
in_features = self.model.roi_heads.box_predictor.cls_score.in_features
self.model.roi_heads.box_predictor = detection.faster_rcnn.FastRCNNPredictor(
in_features, num_categories)
logging.debug('changing num_categories for mask predictor')
in_features_mask = self.model.roi_heads.mask_predictor.conv5_mask.in_channels
self.model.roi_heads.mask_predictor = detection.mask_rcnn.MaskRCNNPredictor(
in_features_mask, 256, num_categories)
def __init__(self, g, in_size, hidden_size, out_size, n_layers,
embedding_size):
super(HeteroRGCN, self).__init__()
# Use trainable node embeddings as featureless inputs.
embed_dict = {
ntype: nn.Parameter(torch.Tensor(g.number_of_nodes(ntype),
in_size))
for ntype in g.ntypes if ntype != 'user'
}
for key, embed in embed_dict.items():
nn.init.xavier_uniform_(embed)
self.embed = nn.ParameterDict(embed_dict)
# create layers
self.layers = nn.Sequential()
self.layers.add_module(
HeteroRGCNLayer(embedding_size, hidden_size, g.etypes))
# hidden layers
for i in range(n_layers - 1):
self.layers.add_module = HeteroRGCNLayer(hidden_size, hidden_size,
g.etypes)
# output layer
self.layers.add(nn.Dense(hidden_size, out_size))
def __init__(
self,
in_features: IntSpaceType,
hidden_multiplier: IntSpaceType,
out_features: IntSpaceType,
act_layer: Callable[[], nn.Module] = nn.GELU,
drop: Optional[float] = None,
):
super(SuperMLPv2, self).__init__()
self._in_features = in_features
self._hidden_multiplier = hidden_multiplier
self._out_features = out_features
self._drop_rate = drop
self._params = nn.ParameterDict({})
self._create_linear("fc1", self.in_features,
int(self.in_features * self.hidden_multiplier))
self._create_linear("fc2",
int(self.in_features * self.hidden_multiplier),
self.out_features)
self.act = act_layer()
self.drop = nn.Dropout(drop or 0.0)
self.reset_parameters()
def __init__(self, graph: dgl.DGLHeteroGraph, input_dim: int, hidden_dim: int, n_heads: int = 4,
dropout: float = 0.2, residual: bool = True):
"""
:param graph: a heterogeneous graph
:param input_dim: int, input dimension
:param hidden_dim: int, hidden dimension
:param n_heads: int, number of attention heads
:param dropout: float, dropout rate
:param residual: boolean, residual connections or not
"""
super(HGConvLayer, self).__init__()
self.input_dim = input_dim
self.hidden_dim = hidden_dim
self.n_heads = n_heads
self.dropout = dropout
self.residual = residual
# hetero conv modules
self.micro_conv = dglnn.HeteroGraphConv({
etype: LSTMConv(dim=input_dim)
for srctype, etype, dsttype in graph.canonical_etypes
})
# different types aggregation module
self.macro_conv = MacroConv(in_feats=hidden_dim * n_heads, out_feats=hidden_dim,
num_heads=n_heads,
dropout=dropout, negative_slope=0.2)
if self.residual:
# residual connection
self.res_fc = nn.ModuleDict()
self.residual_weight = nn.ParameterDict()
for ntype in graph.ntypes:
self.res_fc[ntype] = nn.Linear(input_dim, n_heads * hidden_dim, bias=True)
self.residual_weight[ntype] = nn.Parameter(torch.randn(1))
self.reset_parameters()
def __init__(self, env_func, env_kwargs, state_encoder,
state_dim, action_dim, hidden_dims, activation,
initial_alpha, n_samplers, buffer_capacity,
devices, random_seed=0):
self.devices = itertools.cycle(devices)
self.model_device = next(self.devices)
self.state_dim = state_dim
self.action_dim = action_dim
self.training = True
self.state_encoder = self.STATE_ENCODER_WRAPPER(state_encoder)
self.critic = Critic(state_dim, action_dim, hidden_dims, activation=activation)
self.actor = Actor(state_dim, action_dim, hidden_dims, activation=activation)
self.log_alpha = nn.Parameter(torch.tensor(np.log(initial_alpha), dtype=torch.float32),
requires_grad=True)
self.modules = Container()
self.modules.state_encoder = self.state_encoder
self.modules.critic = self.critic
self.modules.actor = self.actor
self.modules.params = nn.ParameterDict({'log_alpha': self.log_alpha})
self.modules.to(self.model_device)
self.state_encoder.share_memory()
self.actor.share_memory()
self.collector = self.COLLECTOR(env_func=env_func,
env_kwargs=env_kwargs,
state_encoder=self.state_encoder,
actor=self.actor,
n_samplers=n_samplers,
buffer_capacity=buffer_capacity,
devices=self.devices,
random_seed=random_seed)
def __init__(self, f_in, f_out, self_interaction: bool=False, edge_dim: int=0, flavor='skip'):
"""SE(3)-equivariant Graph Conv Layer
Args:
f_in: list of tuples [(multiplicities, type),...]
f_out: list of tuples [(multiplicities, type),...]
self_interaction: include self-interaction in convolution
edge_dim: number of dimensions for edge embedding
flavor: allows ['TFN', 'skip'], where 'skip' adds a skip connection
"""
super().__init__()
self.f_in = f_in
self.f_out = f_out
self.edge_dim = edge_dim
self.self_interaction = self_interaction
self.flavor = flavor
# Neighbor -> center weights
self.kernel_unary = nn.ModuleDict()
for (mi, di) in self.f_in.structure:
for (mo, do) in self.f_out.structure:
self.kernel_unary[f'({di},{do})'] = PairwiseConv(di, mi, do, mo, edge_dim=edge_dim)
# Center -> center weights
self.kernel_self = nn.ParameterDict()
if self_interaction:
assert self.flavor in ['TFN', 'skip']
if self.flavor == 'TFN':
for m_out, d_out in self.f_out.structure:
W = nn.Parameter(torch.randn(1, m_out, m_out) / np.sqrt(m_out))
self.kernel_self[f'{d_out}'] = W
elif self.flavor == 'skip':
for m_in, d_in in self.f_in.structure:
if d_in in self.f_out.degrees:
m_out = self.f_out.structure_dict[d_in]
W = nn.Parameter(torch.randn(1, m_out, m_in) / np.sqrt(m_in))
self.kernel_self[f'{d_in}'] = W
def __init__(self,
G,
in_size,
hidden_size,
out_size,
bias=True,
self_loop=True,
dropout=0.0):
super(HeteroRGCN, self).__init__()
# Use trainable node embeddings as featureless inputs.
embed_dict = {
ntype: nn.Parameter(torch.Tensor(G.number_of_nodes(ntype),
in_size))
for ntype in G.ntypes
}
for key, embed in embed_dict.items():
nn.init.xavier_uniform_(embed)
self.embed = nn.ParameterDict(embed_dict)
# create layers
self.layer1 = HeteroRGCNLayer(in_size, hidden_size, G.etypes, bias,
self_loop, dropout)
self.layer2 = HeteroRGCNLayer(hidden_size, out_size, G.etypes, bias,
self_loop, dropout)
