def forward(self, relations, attributes=None):
batch_size, nr = relations.size()[:2]
assert nr == relations.size(2)
id_shape = list(relations.size()[:-1])
ids = [
random.permutation(2**self.code_length - 1)[:nr] + 1
for i in range(batch_size)
]
ids = np.vstack(ids)
binary_ids = self.binarize_code(ids)
zeros = torch.tensor(np.zeros(binary_ids.shape),
dtype=relations.dtype,
device=relations.device)
binary_ids = torch.tensor(binary_ids,
dtype=relations.dtype,
device=relations.device)
binary_ids2 = torch.cat(meshgrid(binary_ids, dim=1), dim=-1)
if attributes is None:
rels = [binary_ids2, relations]
else:
padding = torch.zeros(*binary_ids2.size()[:-1],
attributes.size(-1),
dtype=relations.dtype,
device=relations.device)
rels = [binary_ids2, padding, relations]
rels = torch.cat(rels, dim=-1)
input_seq = rels.view(batch_size, -1, rels.size(-1))
if attributes is not None:
assert nr == attributes.size(1)
padding = torch.zeros(*binary_ids.size()[:-1],
relations.size(-1),
dtype=relations.dtype,
device=relations.device)
attributes = torch.cat([binary_ids, zeros, attributes, padding],
dim=-1)
input_seq = torch.cat([input_seq, attributes], dim=1)
h0 = torch.zeros(self.num_layers * 2,
batch_size,
self.hidden_size,
dtype=relations.dtype,
device=relations.device)
c0 = torch.zeros(self.num_layers * 2,
batch_size,
self.hidden_size,
dtype=relations.dtype,
device=relations.device)
out, _ = self.lstm(input_seq, (h0, c0))
out = out[:, -1]
if self.feature_dim == 1:
expanded_feature = broadcast(out.unsqueeze(dim=1), 1, nr)
return torch.cat([binary_ids, expanded_feature], dim=-1)
else:
expanded_feature = broadcast(out.unsqueeze(dim=1), 1, nr)
expanded_feature = broadcast(expanded_feature.unsqueeze(dim=1), 1,
nr)
return torch.cat([binary_ids2, expanded_feature], dim=-1)
def randomly_generate_world(nr_blocks, random_order=False, one_stack=False):
"""Randomly generate a blocks world case.
Similar to classical random tree generation, incrementally add new blocks.
for each new block, randomly sample a valid father and stack on its father.
Args:
nr_blocks: The number of blocks in the world.
random_order: Randomly permute the indexes of the blocks if set True.
Or set as a provided order. Leave the raw order unchanged in default.
one_stack: A special case where only one stack of blocks. If True, for each
new node, set its father as the last node.
Returns:
A BlocksWorld instance which is randomly generated.
"""
blocks = [Block(0, None)]
leafs = [blocks[0]]
for i in range(1, nr_blocks + 1):
other = random.choice_list(leafs)
this = Block(i)
this.add_to(other)
if not other.placeable or one_stack:
leafs.remove(other)
blocks.append(this)
leafs.append(this)
order = None
if random_order:
order = random.permutation(len(blocks))
return BlocksWorld(blocks, random_order=order)
def get_order():
raw_order = random.permutation(n - 1)
order = []
for i in range(n - 1):
if i == ground_ind:
order.append(0)
order.append(raw_order[i] + 1)
if ground_ind == n - 1:
order.append(0)
return order
def _restart(self):
"""Restart: Generate a random permutation."""
self._array = random.permutation(self._nr_numbers)
self._set_current_state(self.get_state())
self.optimal = self._calculate_optimal()
def forward(self, relations, attributes=None):
batch_size, nr = relations.size()[:2]
assert nr == relations.size(2)
create_zeros = functools.partial(torch.zeros,
dtype=relations.dtype,
device=relations.device)
id_shape = list(relations.size()[:-1])
ids = [random.permutation(2 ** self.id_dim - 1)[:nr] + 1 \
for i in range(batch_size)]
ids = np.vstack(ids)
binary_ids = self.binarize_code(ids)
zeros = create_zeros(binary_ids.shape)
binary_ids = torch.tensor(binary_ids,
dtype=relations.dtype,
device=relations.device)
binary_ids2 = torch.cat(meshgrid(binary_ids, dim=1), dim=-1)
padded_binary_ids = torch.cat([binary_ids, zeros], dim=-1)
def embed(embed, x):
input_size = x.size()[:-1]
input_channel = x.size(-1)
f = x.view(-1, input_channel)
f = embed(f)
return f.view(*input_size, -1)
if attributes is None:
rels = [binary_ids2, relations]
else:
padding = create_zeros(*binary_ids2.size()[:-1],
attributes.size(-1))
rels = [binary_ids2, padding, relations]
rels = torch.cat(rels, dim=-1)
memory = rels.view(batch_size, -1, rels.size(-1))
if attributes is not None:
assert nr == attributes.size(1)
padding = create_zeros(*padded_binary_ids.size()[:-1],
relations.size(-1))
attributes = torch.cat([padded_binary_ids, attributes, padding],
dim=-1)
memory = torch.cat([memory, attributes], dim=1)
keys = embed(self.key_embed, memory).transpose(1, 2)
values = embed(self.value_embed, memory)
query = padded_binary_ids if self.feature_dim == 1 else binary_ids2
nr_items = nr**self.feature_dim
query = embed(self.query_embed, query).view(batch_size, nr_items, -1)
h0 = create_zeros(batch_size * nr_items, self.hidden_dim)
c0 = create_zeros(batch_size * nr_items, self.hidden_dim)
for i in range(self.queries):
attention = F.softmax(torch.bmm(query, keys), dim=-1)
value = torch.bmm(attention, values)
value = value.view(-1, value.size(-1))
h0, c0 = self.lstm_cell(value, (h0, c0))
query = self.to_query(h0).view(batch_size, nr_items, self.key_dim)
if self.feature_dim == 1:
out = h0.view(batch_size, nr, self.hidden_dim)
else:
out = h0.view(batch_size, nr, nr, self.hidden_dim)
return out
def randomly_generate_family(n, p_marriage=0.8, verbose=False):
"""Randomly generate family trees.
Mimic the process of families growing using a timeline. Each time a new person
is created, randomly sample the gender and parents (could be none, indicating
not included in the family tree) of the person. Also maintain lists of singles
of each gender. With probability $p_marrige, randomly pick two from each list
to be married. Finally randomly permute the order of people.
Args:
n: The number of people in the family tree.
p_marriage: The probability of marriage happens each time.
verbose: print the marriage and child born process if verbose=True.
Returns:
A family tree instance of $n people.
"""
assert n > 0
ids = list(random.permutation(n))
single_m = []
single_w = []
couples = [None]
# The relations are: husband, wife, father, mother, son, daughter
rel = np.zeros((n, n, 6))
fathers = [None for i in range(n)]
mothers = [None for i in range(n)]
def add_couple(man, woman):
"""Add a couple relation among (man, woman)."""
couples.append((man, woman))
rel[woman, man, 0] = 1 # husband
rel[man, woman, 1] = 1 # wife
if verbose:
print('couple', man, woman)
def add_child(parents, child, gender):
"""Add a child relation between parents and the child according to gender."""
father, mother = parents
fathers[child] = father
mothers[child] = mother
rel[child, father, 2] = 1 # father
rel[child, mother, 3] = 1 # mother
if gender == 0: # son
rel[father, child, 4] = 1
rel[mother, child, 4] = 1
else: # daughter
rel[father, child, 5] = 1
rel[mother, child, 5] = 1
if verbose:
print('child', father, mother, child, gender)
def check_relations(man, woman):
"""Disable marriage between cousins."""
if fathers[man] is None or fathers[woman] is None:
return True
if fathers[man] == fathers[woman]:
return False
def same_parent(x, y):
return fathers[x] is not None and fathers[
y] is not None and fathers[x] == fathers[y]
for x in [fathers[man], mothers[man]]:
for y in [fathers[woman], mothers[woman]]:
if same_parent(man, y) or same_parent(woman, x) or same_parent(
x, y):
return False
return True
while ids:
x = ids.pop()
gender = random.randint(2)
parents = random.choice(couples)
if gender == 0:
single_m.append(x)
else:
single_w.append(x)
if parents is not None:
add_child(parents, x, gender)
if random.rand() < p_marriage and len(single_m) > 0 and len(
single_w) > 0:
mi = random.randint(len(single_m))
wi = random.randint(len(single_w))
man = single_m[mi]
woman = single_w[wi]
if check_relations(man, woman):
add_couple(man, woman)
del single_m[mi]
del single_w[wi]
return Family(n, rel)