def __init__(self, dset, demo_hidden, pooling):
super(DemoEncoder, self).__init__()
# Image encoder
import torchvision
resnet = torchvision.models.resnet18(pretrained=True)
resnet = list(resnet.children())[:-1]
self.resnet = nn.Sequential(*resnet)
self.seg_map = SegmapCNN()
self.seg_feat2hidden = nn.Sequential()
self.seg_feat2hidden.add_module(
'fc_block1', fc_block(
4160, demo_hidden, False, nn.ReLU))
feat2hidden = nn.Sequential()
feat2hidden.add_module(
'fc_block1',
fc_block(512, demo_hidden, False, nn.ReLU))
self.feat2hidden = feat2hidden
self.combine = nn.Sequential()
self.combine.add_module(
'fc_block1',
fc_block(
2 * demo_hidden,
demo_hidden,
False,
nn.ReLU))
self.pooling = pooling
def __init__(
self,
args,
dset):
super(PredicateClassifier, self).__init__()
self.num_goal_predicates = dset.num_goal_predicates
self.max_possible_count = dset.max_goal_length
hidden_size = args.demo_hidden
print('hidden_size', hidden_size)
if args.dropout==0:
print('dropout', args.dropout)
classifier = nn.Sequential()
classifier.add_module('fc_block1', fc_block(hidden_size, hidden_size, False, nn.Tanh))
classifier.add_module('fc_block2', fc_block(hidden_size, self.max_possible_count*self.num_goal_predicates, False, None))
else:
print('dropout not 0', args.dropout)
classifier = nn.Sequential()
classifier.add_module('fc_block1', fc_block(hidden_size, hidden_size, False, nn.Tanh))
classifier.add_module('dropout', nn.Dropout(args.dropout))
classifier.add_module('fc_block2', fc_block(hidden_size, self.max_possible_count*self.num_goal_predicates, False, None))
self.classifier = classifier
def __init__(self, dset, embedding_dim, hidden):
super(ProgramEncoder, self).__init__()
num_actions = dset.num_actions
num_objects = dset.num_objects
num_indexes = dset.num_indexes
# Encode sketch
action_embedding = nn.Embedding(num_actions, embedding_dim)
object_embedding = nn.Embedding(num_objects, embedding_dim)
index_embedding = nn.Embedding(num_indexes, embedding_dim)
object_index2instance = nn.Sequential()
object_index2instance.add_module(
'fc1', fc_block(2 * embedding_dim, embedding_dim, False, nn.Tanh))
object_index2instance.add_module(
'fc2', fc_block(embedding_dim, embedding_dim, False, nn.Tanh))
self.object_index2instance = object_index2instance
atomic_action2emb = nn.Sequential()
atomic_action2emb.add_module(
'fc_block1', fc_block(embedding_dim + 2 * embedding_dim, hidden, False, nn.Tanh))
atomic_action2emb.add_module(
'fc_block2', fc_block(hidden, hidden, False, nn.Tanh))
self.atomic_action2emb = atomic_action2emb
sketch_encoding = nn.GRU(hidden, hidden, 2, batch_first=True)
self.object_dict = dset.object_dict
self.action_embedding = action_embedding
self.object_embedding = object_embedding
self.index_embedding = index_embedding
self.sketch_encoding = sketch_encoding
def __init__(self, dset, embedding_dim, desc_hidden, word_embedding):
super(WordEncoder, self).__init__()
emb2hidden = nn.Sequential()
emb2hidden.add_module('fc_block1', fc_block(embedding_dim, desc_hidden, True, nn.Tanh))
emb2hidden.add_module('fc_block2', fc_block(desc_hidden, desc_hidden, True, nn.Tanh))
self.emb2hidden = emb2hidden
word_encoding = nn.GRU(desc_hidden, desc_hidden, 2)
self.object_dict = dset.object_dict
self.word_embedding = word_embedding
self.word_encoding = word_encoding
def __init__(self, n_timesteps, n_edge_types, graph_hidden, embedding_dim,
hidden):
super(VanillaGraphEncoder, self).__init__()
layernorm = True
self.n_timesteps = n_timesteps
self.n_edge_types = n_edge_types
self.embedding_dim = embedding_dim
self.input_dim = n_edge_types + embedding_dim
self.graph_hidden = graph_hidden
node_init2hidden = nn.Sequential()
node_init2hidden.add_module(
'fc1', fc_block(3 * embedding_dim, graph_hidden, False, nn.Tanh))
node_init2hidden.add_module(
'fc2', fc_block(graph_hidden, graph_hidden, False, nn.Tanh))
for i in range(n_edge_types):
hidden2message_in = fc_block(graph_hidden, graph_hidden, False,
nn.Tanh)
self.add_module("hidden2message_in_{}".format(i),
hidden2message_in)
hidden2message_out = fc_block(graph_hidden, graph_hidden, False,
nn.Tanh)
self.add_module("hidden2message_out_{}".format(i),
hidden2message_out)
if layernorm:
self.gru_cell = LayerNormGRUCell
else:
self.gru_cell = nn.GRUCell
propagator = self.gru_cell(input_size=2 * n_edge_types * graph_hidden,
hidden_size=graph_hidden)
self.node_init2hidden = node_init2hidden
self.hidden2message_in = AttrProxy(self, "hidden2message_in_")
self.hidden2message_out = AttrProxy(self, "hidden2message_out_")
self.propagator = propagator
def __init__(self, n_edge_types, n_touch, graph_hidden, embedding_dim,
hidden):
super(GRUActionGraphEncoder,
self).__init__(0, n_edge_types, graph_hidden, embedding_dim,
hidden)
self.n_touch = n_touch
action2hidden = nn.Sequential()
action2hidden.add_module(
'fc1',
fc_block(embedding_dim + n_touch, graph_hidden, False, nn.Tanh))
action2hidden.add_module(
'fc2', fc_block(graph_hidden, graph_hidden, False, nn.Tanh))
temporal_propagator = self.gru_cell(input_size=graph_hidden,
hidden_size=graph_hidden)
self.temporal_propagator = temporal_propagator
self.action2hidden = action2hidden
def __init__(self, n_edge_types, n_touch, graph_hidden, embedding_dim,
hidden):
super(FCActionGraphEncoder,
self).__init__(0, n_edge_types, graph_hidden, embedding_dim,
hidden)
self.n_touch = n_touch
action2hidden = nn.Sequential()
action2hidden.add_module(
'fc1',
fc_block(embedding_dim + n_touch, graph_hidden, False, nn.Tanh))
action2hidden.add_module(
'fc2', fc_block(graph_hidden, graph_hidden, False, nn.Tanh))
compute_residual = nn.Sequential()
compute_residual.add_module(
'fc1', fc_block(2 * graph_hidden, graph_hidden, False, nn.Tanh))
compute_residual.add_module(
'fc2', fc_block(graph_hidden, graph_hidden, False, nn.Tanh))
self.compute_residual = compute_residual
self.action2hidden = action2hidden
def __init__(self, args, dset, pooling):
super(ActionDemoEncoder, self).__init__()
# Image encoder
import torchvision
hidden_size = args.demo_hidden
len_action_predicates = len(dset.action_predicates)
self.action_embed = nn.Embedding(len_action_predicates, hidden_size)
feat2hidden = nn.Sequential()
feat2hidden.add_module(
'fc_block1', fc_block(hidden_size, hidden_size, False, nn.ReLU))
self.feat2hidden = feat2hidden
self.pooling = pooling
if 'lstm' in self.pooling:
self.lstm = nn.LSTM(hidden_size, hidden_size)
def __init__(
self,
dset,
desc_hidden,
embedding_dim,
hidden,
max_words):
super(ProgramDecoder, self).__init__()
layernorm = True
num_actions = dset.num_actions
num_objects = dset.num_objects
# Encode sketch embedding
input_emb_encoder = nn.Sequential()
input_emb_encoder.add_module(
'fc_block1', fc_block(desc_hidden, hidden, False, nn.Tanh))
self.input_emb_encoder = input_emb_encoder
# Encode input words
action_embedding = nn.Embedding(num_actions, embedding_dim)
object_embedding = nn.Embedding(num_objects, embedding_dim)
self.action_embedding = action_embedding
self.object_embedding = object_embedding
atomic_action2emb = nn.Sequential()
atomic_action2emb.add_module(
'fc_block1', fc_block(3 * embedding_dim, hidden, False, nn.Tanh))
atomic_action2emb.add_module(
'fc_block2', fc_block(hidden, hidden, False, nn.Tanh))
self.atomic_action2emb = atomic_action2emb
if layernorm:
gru_cell = LayerNormGRUCell
else:
gru_cell = nn.GRUCell
decode_gru_cell = gru_cell(2 * hidden, hidden)
gru2hidden = nn.Sequential()
gru2hidden.add_module(
'fc_block1', fc_block(hidden, hidden, False, nn.Tanh))
gru2hidden.add_module(
'fc_block2', fc_block(hidden, hidden, False, nn.Tanh))
hidden2action = nn.Sequential()
hidden2action.add_module(
'fc_block1', fc_block(hidden, hidden, False, nn.Tanh))
hidden2action.add_module(
'fc_block2', fc_block(hidden, num_actions, False, None))
hidden2object1 = nn.Sequential()
hidden2object1.add_module(
'fc_block1', fc_block(hidden + embedding_dim, hidden, False, nn.Tanh))
hidden2object1.add_module(
'fc_block2', fc_block(hidden, num_objects, False, None))
hidden2object2 = nn.Sequential()
hidden2object2.add_module(
'fc_block1', fc_block(hidden + 2 * embedding_dim, hidden, False, nn.Tanh))
hidden2object2.add_module(
'fc_block2', fc_block(hidden, num_objects, False, None))
self.embedding_dim = embedding_dim
self.decode_gru_cell = decode_gru_cell
self.gru2hidden = gru2hidden
self.hidden2action = hidden2action
self.hidden2object1 = hidden2object1
self.hidden2object2 = hidden2object2
self.max_words = max_words
def __init__(
self,
dset,
sketch_embedding_dim,
embedding_dim,
hidden,
max_words,
graph_state_dim):
super(ProgramClassifier, self).__init__()
layernorm = True
num_actions = dset.num_actions
num_objects = dset.num_objects
# Encode sketch embedding
sketch_emb_encoder = nn.Sequential()
sketch_emb_encoder.add_module(
'fc_block1', fc_block(sketch_embedding_dim, hidden, False, nn.Tanh))
# Encode input words
action_embedding = nn.Embedding(num_actions, embedding_dim)
object_embedding = nn.Embedding(num_objects, embedding_dim)
atomic_action2emb = nn.Sequential()
atomic_action2emb.add_module(
'fc_block1', fc_block(embedding_dim + 2 * graph_state_dim, hidden, False, nn.Tanh))
atomic_action2emb.add_module(
'fc_block2', fc_block(hidden, hidden, False, nn.Tanh))
if layernorm:
gru_cell = LayerNormGRUCell
else:
gru_cell = nn.GRUCell
# gru
decode_gru_cell = gru_cell(2 * hidden, hidden)
# decode words
gru2hidden = nn.Sequential()
gru2hidden.add_module(
'fc_block1', fc_block(hidden, hidden, False, nn.Tanh))
gru2hidden.add_module(
'fc_block2', fc_block(hidden, hidden, False, nn.Tanh))
hidden2object1_logits = nn.Sequential()
hidden2object1_logits.add_module(
'fc_block1', fc_block(hidden + graph_state_dim, hidden, False, nn.Tanh))
hidden2object1_logits.add_module(
'fc_block2', fc_block(hidden, 1, False, None))
hidden2action_logits = nn.Sequential()
hidden2action_logits.add_module(
'fc_block1', fc_block(hidden + graph_state_dim, hidden, False, nn.Tanh))
hidden2action_logits.add_module(
'fc_block2', fc_block(hidden, num_actions, False, None))
hidden2object2_logits = nn.Sequential()
hidden2object2_logits.add_module(
'fc_block1', fc_block(hidden + 2 * graph_state_dim + embedding_dim, hidden, False, nn.Tanh))
hidden2object2_logits.add_module(
'fc_block2', fc_block(hidden, 1, False, None))
self.embedding_dim = embedding_dim
self.sketch_emb_encoder = sketch_emb_encoder
self.action_embedding = action_embedding
self.object_embedding = object_embedding
self.atomic_action2emb = atomic_action2emb
self.decode_gru_cell = decode_gru_cell
self.gru2hidden = gru2hidden
self.hidden2object1_logits = hidden2object1_logits
self.hidden2action_logits = hidden2action_logits
self.hidden2object2_logits = hidden2object2_logits
self.num_actions = num_actions
self.max_words = max_words
def __init__(
self,
dset,
env_graph_encoder,
env_graph_helper,
sketch_embedding_dim,
embedding_dim,
hidden,
max_words,
graph_evolve,
interaction_grpah_helper=None,
interaction_graph_encoder=None):
super(ProgramGraphClassifier, self).__init__()
if graph_evolve:
self.interaction_grpah_helper = interaction_grpah_helper
self.interaction_graph_encoder = interaction_graph_encoder
env_graph2interactoin_graph = nn.Sequential()
env_graph2interactoin_graph.add_module(
'fc1',
fc_block(
env_graph_encoder.graph_hidden,
interaction_graph_encoder.graph_hidden,
False,
nn.Tanh))
self.env_graph2interactoin_graph = env_graph2interactoin_graph
graph_state_dim = interaction_graph_encoder.graph_hidden
else:
graph_state_dim = env_graph_encoder.graph_hidden
program_classifier = ProgramClassifier(
dset, sketch_embedding_dim, embedding_dim, hidden, max_words, graph_state_dim)
category_embedding = nn.Embedding(
env_graph_helper.n_node_category, embedding_dim)
node_state2_hidden = fc_block(
env_graph_helper.n_node_state,
embedding_dim,
False,
nn.Tanh)
self.hidden = hidden
self.action_object_constraints = dset.compatibility_matrix
self.object_dict = dset.object_dict
self.env_graph_helper = env_graph_helper
self.env_graph_encoder = env_graph_encoder
self.program_classifier = program_classifier
self.category_embedding = category_embedding
self.node_state2_hidden = node_state2_hidden
self.graph_evolve = graph_evolve
graph_feat2hx = nn.Sequential()
graph_feat2hx.add_module(
'fc1',
fc_block(
graph_state_dim,
hidden,
False,
nn.Tanh))
graph_feat2hx.add_module(
'fc2', fc_block(
hidden, hidden, False, nn.Tanh))
self.graph_feat2hx = graph_feat2hx