def increase_input_size(self, new_pattern_dim, new_graph_dim):
assert new_pattern_dim >= self.pattern_dim and new_graph_dim >= self.graph_dim
if new_pattern_dim != self.pattern_dim:
new_p_layer = extend_dimensions(self.p_layer, new_input_dim=new_pattern_dim, upper=False)
del self.p_layer
self.p_layer = new_p_layer
if new_graph_dim != self.graph_dim:
new_g_layer = extend_dimensions(self.g_layer, new_input_dim=new_graph_dim, upper=False)
del self.g_layer
self.g_layer = new_g_layer
self.pattern_dim = new_pattern_dim
self.graph_dim = new_graph_dim
def increase_input_size(self, new_query_dim, new_key_dim, new_value_dim, new_mem_dim):
if self.mem_init.endswith("attn"):
self.m_layer.increase_input_size(new_query_dim, new_key_dim, new_value_dim)
else:
new_m_layer = extend_dimensions(self.m_layer, new_input_dim=new_key_dim, new_output_dim=new_mem_dim, upper=False)
del self.m_layer
self.m_layer = new_m_layer
self.attn.increase_input_size(new_query_dim, new_mem_dim, new_mem_dim)
def increase_input_size(self, new_input_dim):
assert new_input_dim >= self.input_dim
if new_input_dim != self.input_dim:
new_emb_layer = extend_dimensions(self.emb_layer,
new_input_dim=new_input_dim,
upper=False)
del self.emb_layer
self.emb_layer = new_emb_layer
self.input_dim = new_input_dim
def increase_input_size(self, new_query_dim, new_key_dim, new_value_dim):
assert new_query_dim >= self.query_dim and new_key_dim >= self.key_dim and new_value_dim >= new_value_dim
if new_query_dim != self.query_dim:
new_q_net = extend_dimensions(self.q_net, new_input_dim=new_query_dim, upper=False)
del self.q_net
self.q_net = new_q_net
if self.pre_lnorm:
new_q_layer_norm = extend_dimensions(self.q_layer_norm, new_input_dim=new_query_dim, upper=False)
del self.q_layer_norm
self.q_layer_norm = new_q_layer_norm
new_o_net = extend_dimensions(self.o_net, new_output_dim=new_query_dim, upper=False)
del self.o_net
self.o_net = new_o_net
if self.post_lnorm:
new_o_layer_norm = extend_dimensions(self.o_layer_norm, new_input_dim=new_query_dim, upper=False)
del self.o_layer_norm
self.o_layer_norm = new_o_layer_norm
new_g_net = nn.Linear(2*new_query_dim, new_query_dim)
with torch.no_grad():
nn.init.zeros_(new_g_net.weight)
nn.init.zeros_(new_g_net.bias)
new_g_net.weight[-self.query_dim:, new_query_dim-self.query_dim:new_query_dim].data.copy_(self.g_net.weight[:, :self.query_dim])
new_g_net.weight[-self.query_dim:, -self.query_dim:].data.copy_(self.g_net.weight[:, -self.query_dim:])
new_g_net.bias[new_query_dim-self.query_dim:new_query_dim].data.copy_(self.g_net.bias[:self.query_dim])
new_g_net.bias[-self.query_dim:].data.copy_(self.g_net.bias[-self.query_dim:])
del self.g_net
self.g_net = new_g_net
if new_key_dim != self.key_dim:
new_k_net = extend_dimensions(self.k_net, new_input_dim=new_key_dim, upper=False)
del self.k_net
self.k_net = new_k_net
if self.pre_lnorm:
new_k_layer_norm = extend_dimensions(self.k_layer_norm, new_input_dim=new_key_dim, upper=False)
del self.k_layer_norm
self.k_layer_norm = new_k_layer_norm
if new_value_dim != self.value_dim:
new_v_net = extend_dimensions(self.v_net, new_input_dim=new_value_dim, upper=False)
del self.v_net
self.v_net = new_v_net
if self.pre_lnorm:
new_v_layer_norm = extend_dimensions(self.v_layer_norm, new_input_dim=new_value_dim, upper=False)
del self.v_layer_norm
self.v_layer_norm = new_v_layer_norm
self.query_dim = new_query_dim
self.key_dim = new_key_dim
self.value_dim = new_value_dim
def increase_input_size(self, new_pattern_dim, new_graph_dim):
assert new_pattern_dim >= self.pattern_dim and new_graph_dim >= self.graph_dim
if new_pattern_dim != self.pattern_dim:
new_p_layer = extend_dimensions(self.p_layer, new_input_dim=new_pattern_dim, upper=False)
del self.p_layer
self.p_layer = new_p_layer
if new_graph_dim != self.graph_dim:
new_g_layer = extend_dimensions(self.g_layer, new_input_dim=new_graph_dim, upper=False)
del self.g_layer
self.g_layer = new_g_layer
if self.mem_init.endswith("attn"):
self.m_layer.increase_input_size(self.hidden_dim, new_graph_dim, new_graph_dim)
else:
new_m_layer = extend_dimensions(self.m_layer, new_input_dim=new_graph_dim, upper=False)
del self.m_layer
self.m_layer = new_m_layer
self.p_attn.increase_input_size(self.hidden_dim, new_pattern_dim, new_pattern_dim)
self.g_attn.increase_input_size(self.hidden_dim, new_graph_dim, new_graph_dim)
self.pattern_dim = new_pattern_dim
self.graph_dim = new_graph_dim
def increase_input_size(self, new_query_dim, new_key_dim, new_value_dim):
assert new_query_dim >= self.query_dim and new_key_dim >= self.key_dim and new_value_dim >= new_value_dim
if new_query_dim != self.query_dim:
new_q_net = extend_dimensions(self.q_net, new_input_dim=new_query_dim, upper=False)
del self.q_net
self.q_net = new_q_net
if self.pre_lnorm:
new_q_layer_norm = extend_dimensions(self.q_layer_norm, new_input_dim=new_query_dim, upper=False)
del self.q_layer_norm
self.q_layer_norm = new_q_layer_norm
new_o_net = extend_dimensions(self.o_net, new_output_dim=new_query_dim, upper=False)
del self.o_net
self.o_net = new_o_net
if self.post_lnorm:
new_o_layer_norm = extend_dimensions(self.o_layer_norm, new_input_dim=new_query_dim, upper=False)
del self.o_layer_norm
self.o_layer_norm = new_o_layer_norm
if new_key_dim != self.key_dim:
new_k_net = extend_dimensions(self.k_net, new_input_dim=new_key_dim, upper=False)
del self.k_net
self.k_net = new_k_net
if self.pre_lnorm:
new_k_layer_norm = extend_dimensions(self.k_layer_norm, new_input_dim=new_key_dim, upper=False)
del self.k_layer_norm
self.k_layer_norm = new_k_layer_norm
if new_value_dim != self.value_dim:
new_v_net = extend_dimensions(self.v_net, new_input_dim=new_value_dim, upper=False)
del self.v_net
self.v_net = new_v_net
if self.pre_lnorm:
new_v_layer_norm = extend_dimensions(self.v_layer_norm, new_input_dim=new_value_dim, upper=False)
del self.v_layer_norm
self.v_layer_norm = new_v_layer_norm
self.query_dim = new_query_dim
self.key_dim = new_key_dim
self.value_dim = new_value_dim