def test_transformer_parameters_to_attention(self):
with self.assertRaises(ValueError):
transformer = TransformerEncoderBuilder.from_kwargs(
attention_type="test-attention").get()
transformer = TransformerEncoderBuilder.from_kwargs(
attention_type="test-attention", n_heads=8,
query_dimensions=64).get()
def test_simple_build(self):
transformer = TransformerEncoderBuilder().get()
builder = TransformerEncoderBuilder()
builder.n_layers = 1
builder.n_heads = 4
builder.attention_type = "linear"
transformer = builder.get()
with self.assertRaises(ValueError):
builder = TransformerEncoderBuilder()
builder.attention_type = "whatever"
def __init__(self, config, output_attentions=False, keep_multihead_output=False):
super(TransformerEncoder, self).__init__()
self.output_attentions = output_attentions
self.pre_layer_norm = config.pre_layer_norm
builder = TransformerEncoderBuilder.from_kwargs(
n_layers=config.num_hidden_layers,
n_heads=config.num_attention_heads,
feed_forward_dimensions=config.intermediate_size,
query_dimensions=int(config.hidden_size / config.num_attention_heads),
value_dimensions=int(config.hidden_size / config.num_attention_heads),
dropout=config.hidden_dropout_prob
)
if config.softmax_temp:
builder.attention.softmax_temp = config.softmax_temp
builder.attention.attention_dropout = config.attention_probs_dropout_prob
builder.attention.clusters = config.clusters
builder.attention.bits = config.bits
builder.attention.hash_bias = config.hash_bias
builder.attention.iterations = config.iterations
builder.attention.topk = config.topk
builder.attention.local_context = config.local_context
builder.attention.length_limit = config.length_limit
attention_type = config.attention_type
if attention_type == "improved-clustered":
attention_type = "conditional-full:improved-clustered"
builder.attention_type = attention_type
self.transformer = builder.get()
def test_longformer(self):
config = LongformerConfig()
config.attention_mode = "n2"
config.attention_window = [256] * 12
config.attention_dilation = [1] * 12
longformer = Longformer(config)
encoder = TransformerEncoderBuilder.from_kwargs(
n_layers=12,
n_heads=12,
query_dimensions=64,
value_dimensions=64,
feed_forward_dimensions=3072,
attention_type="full",
final_normalization=False,
activation="gelu").get()
longformer.eval()
encoder.eval()
# Before the weight copy they should be different
x = torch.rand(3, 10, 768)
o1 = longformer.encoder(x, head_mask=[None] * 12)[0]
o2 = encoder(x)
self.assertGreater(torch.abs(o1 - o2).max().item(), 1)
# And after the copy they should be exactly the same
encoder.load_state_dict(LongformerMapper().map(
longformer.encoder.state_dict()))
o1 = longformer.encoder(x, head_mask=[None] * 12)[0]
o2 = encoder(x)
self.assertLess(torch.abs(o1 - o2).max().item(), 1e-4)
def test_huggin_bert(self):
bert = BertModel(BertConfig())
encoder = TransformerEncoderBuilder.from_kwargs(
n_layers=12,
n_heads=12,
query_dimensions=64,
value_dimensions=64,
feed_forward_dimensions=3072,
attention_type="full",
final_normalization=False,
activation="gelu").get()
bert.eval()
encoder.eval()
# Before the weight copy they should be different
x = torch.rand(3, 10, 768)
o1 = bert.encoder(x, head_mask=[None] * 12)[0]
o2 = encoder(x)
self.assertGreater(torch.abs(o1 - o2).max().item(), 1)
# And after the copy they should be exactly the same
encoder.load_state_dict(HugginfaceBertEncoderMapper().map(
bert.encoder.state_dict()))
o1 = bert.encoder(x, head_mask=[None] * 12)[0]
o2 = encoder(x)
self.assertLess(torch.abs(o1 - o2).max().item(), 1e-4)
def __init__(self, num_feats, output_feats, lstm_layers, n_layers,
n_heads, hidden_dim, ff_dim, tf_depth=3, dropout=0.15):
super(LSTUT, self).__init__()
self.num_feats = num_feats
self.output_feats = output_feats
self.lstm_layers = lstm_layers
self.n_layers = n_layers
self.n_heads = n_heads
self.hidden_dim = hidden_dim
self.ff_dim = ff_dim
self.tf_depth = tf_depth
self.d_model = self.hidden_dim * self.n_heads
encoder_builder = TransformerEncoderBuilder.from_kwargs(
n_layers=self.n_layers,
n_heads=self.n_heads,
query_dimensions=self.hidden_dim,
value_dimensions=self.hidden_dim,
feed_forward_dimensions=self.ff_dim,
attention_type='linear',
dropout=dropout
)
self.initial_ff = nn.Linear(self.num_feats, self.d_model)
self.lstm1 = nn.LSTM(self.d_model, self.d_model // 2, self.lstm_layers, batch_first=True, bidirectional=True)
self.encoder = encoder_builder.get()
self.lstm2 = nn.LSTM(self.d_model, self.d_model // 2, self.lstm_layers, batch_first=True, bidirectional=True)
self.final_ff = nn.Linear(self.d_model, self.output_feats)
def test_attention_composition(self):
transformer = TransformerEncoderBuilder.from_kwargs(
attention_type="conditional-full:improved-clustered",
attention_dropout=0.1,
softmax_temp=0.125,
clusters=256,
bits=32,
topk=32,
length_limit=512).get()
with self.assertRaises(TypeError):
transformer = TransformerEncoderBuilder.from_kwargs(
attention_type="conditional-full",
attention_dropout=0.1,
softmax_temp=0.125,
length_limit=512).get()
def __init__(
self,
attention_type,
out_channels,
num_layers=6,
nhead=8,
d_ffn=1024,
dropout=0,
activation="relu",
reformer_bucket_size=32,
):
super(FastTransformerBlock, self).__init__()
from fast_transformers.builders import TransformerEncoderBuilder
builder = TransformerEncoderBuilder.from_kwargs(
attention_type=attention_type,
n_layers=num_layers,
n_heads=nhead,
feed_forward_dimensions=d_ffn,
query_dimensions=out_channels // nhead,
value_dimensions=out_channels // nhead,
dropout=dropout,
attention_dropout=dropout,
chunk_size=reformer_bucket_size,
)
self.mdl = builder.get()
self.attention_type = attention_type
self.reformer_bucket_size = reformer_bucket_size
def __init__(
self,
file_re,
input_dimensions=5,
# project_dimension=128,
n_layers=8,
n_heads=8,
query_dimensions=64,
value_dimensions=64,
feed_forward_dimensions=1024,
attention_type='full',
num_workers=1,
batch_size=2,
lr=1e-7,
seq_len=1000,
seed=100,
**kwargs):
super(TimeSeriesTransformer, self).__init__()
builder = TransformerEncoderBuilder.from_kwargs(
n_layers=n_layers,
n_heads=n_heads,
query_dimensions=query_dimensions,
value_dimensions=value_dimensions,
feed_forward_dimensions=feed_forward_dimensions)
# Build a transformer with softmax attention
builder.attention_type = attention_type
self.transformer = builder.get()
project_dimension = query_dimensions * n_heads
self.project_dimension = project_dimension
# self.stock_num = stock_num
# self.stock_mapping_fn = stock_mapping_fn
# self.model_projection = nn.Conv1d(input_dimensions, project_dimension, 1)
self.model_projection = nn.Linear(input_dimensions, project_dimension)
self.positional_encoder = PositionalEncoding(project_dimension)
self.seq_len = seq_len
self.loss = MaskedRMSLE()
self.file_re = file_re
self.batch_size = batch_size
self.metric_object = MaskedAPE()
self.output_projection = nn.Linear(n_heads * value_dimensions, 1)
# self.output_projection = nn.Conv1d(n_heads * value_dimensions, 1, 1)
self.filenames = glob.glob(self.file_re)
self.lr = lr
self.seed = seed
np.random.seed(seed)
# self.split_date = split_date
# self.end_date = end_date
# random.shuffle(self.filenames)
# self.training_files = self.filenames
# self.valid_files = glob.glob(valid_file_re)
# if not self.training_files:
# # if not self.training_files or (not self.valid_files):
# raise ValueError(f"input file train {self.training_files} is empty")
self.num_workers = num_workers
def __init__(self,
d_model,
sequence_length,
n_classes,
attention_type="full",
n_layers=4,
n_heads=4,
d_query=32,
dropout=0.1,
softmax_temp=None,
attention_dropout=0.1,
bits=32,
rounds=4,
chunk_size=32,
masked=True):
super(SequencePredictor, self).__init__()
self.pos_embedding = PositionalEncoding(d_model // 2,
max_len=sequence_length)
self.value_embedding = torch.nn.Embedding(n_classes + 1, d_model // 2)
self.builder_dict = OrderedDict({
"attention_type":
attention_type,
"n_layers":
n_layers,
"n_heads":
n_heads,
"feed_forward_dimensions":
n_heads * d_query * 4,
"query_dimensions":
d_query,
"value_dimensions":
d_query,
"dropout":
dropout,
"softmax_temp":
softmax_temp,
"attention_dropout":
attention_dropout,
"bits":
bits,
"rounds":
rounds,
"chunk_size":
chunk_size,
"masked":
masked
})
self.transformer = TransformerEncoderBuilder.from_dictionary(
self.builder_dict, strict=True).get()
hidden_size = n_heads * d_query
self.predictor = torch.nn.Linear(hidden_size, n_classes + 1)
def test_attention_matrix(self):
A = []
def store_attention(event):
A.append(event.attention_matrix)
# default transformer is 4 layers 4 heads
transformer = TransformerEncoderBuilder().get()
x = transformer(torch.rand(1, 100, 64*4))
self.assertEqual(len(A), 0)
EventDispatcher.get().listen(AttentionEvent, store_attention)
x = transformer(torch.rand(1, 100, 64*4))
self.assertEqual(len(A), 4)
def __init__(self, input_feats, output_feats, n_layers, n_heads, hidden_dim, ff_dim, tf_depth=3, dropout=0.15):
super(TransformerEncoderDecoder, self).__init__()
self.input_feats = input_feats
self.output_feats = output_feats
self.n_layers = n_layers
self.n_heads = n_heads
self.hidden_dim = hidden_dim
self.ff_dim = ff_dim
self.d_model = hidden_dim * n_heads
self.tf_depth = tf_depth
self.pe = PositionalEncoding(self.d_model, dropout=dropout, max_len=4096)
self.A = nn.GELU()
encoder_builder = TransformerEncoderBuilder.from_kwargs(
n_layers=n_layers,
n_heads=n_heads,
query_dimensions=hidden_dim,
value_dimensions=hidden_dim,
feed_forward_dimensions=ff_dim,
attention_type='linear',
dropout=dropout
)
decoder_builder = TransformerDecoderBuilder.from_kwargs(
n_layers=n_layers,
n_heads=n_heads,
query_dimensions=hidden_dim,
value_dimensions=hidden_dim,
feed_forward_dimensions=ff_dim,
cross_attention_type='linear',
self_attention_type='causal-linear',
dropout=dropout
)
self.encoder = encoder_builder.get()
self.decoder = decoder_builder.get()
self.src_embed = nn.Linear(self.input_feats, self.d_model)
self.tgt_embed = nn.Linear(self.output_feats, self.d_model)
self.final_ff = nn.Linear(self.d_model, self.output_feats)
self.tgt_mask = fast_transformers.masking.TriangularCausalMask(N=12)
def __init__(self,
d_model,
coords_target, flatten_order_target,
attention_type="linear",
n_layers=4,
n_heads=4,
d_query=32,
dropout=0.1,
attention_dropout=0.1,
d_conv=8):
super(TRecOnlyFBP, self).__init__()
self.fbp_fourier_coefficient_embedding = torch.nn.Linear(2, d_model // 2)
self.pos_embedding_target = PositionalEncoding2D(d_model // 2, coords=coords_target,
flatten_order=flatten_order_target)
self.encoder = TransformerEncoderBuilder.from_kwargs(
attention_type=attention_type,
n_layers=n_layers,
n_heads=n_heads,
feed_forward_dimensions=n_heads * d_query * 4,
query_dimensions=d_query,
value_dimensions=d_query,
dropout=dropout,
attention_dropout=attention_dropout
).get()
self.predictor_amp = torch.nn.Linear(
n_heads * d_query,
1
)
self.predictor_phase = torch.nn.Linear(
n_heads * d_query,
1
)
self.conv_block = torch.nn.Sequential(
torch.nn.Conv2d(1, d_conv, kernel_size=3, stride=1, padding=1),
torch.nn.ReLU(),
torch.nn.BatchNorm2d(d_conv),
torch.nn.Conv2d(d_conv, 1, kernel_size=1, stride=1, padding=0)
)
def test_builder_factory_methods(self):
builder = TransformerEncoderBuilder.from_kwargs(
n_layers=1, n_heads=4, attention_type="linear")
with self.assertRaises(ValueError):
TransformerEncoderBuilder.from_kwargs(foobar=1)
TransformerEncoderBuilder.from_kwargs(foobar=1, strict=False)
parser = argparse.ArgumentParser()
parser.add_argument("--n_layers", type=int)
parser.add_argument("--n_heads", type=int)
args = parser.parse_args(["--n_heads", "42"])
builder = TransformerEncoderBuilder.from_namespace(args)
self.assertEqual(builder.n_heads, 42)
self.assertTrue(builder.n_layers is None)
def test_mapping(self):
t1 = nn.TransformerEncoder(
nn.TransformerEncoderLayer(128, 4, dim_feedforward=256), 4)
t2 = TransformerEncoderBuilder.from_kwargs(
n_layers=4,
n_heads=4,
query_dimensions=128 // 4,
value_dimensions=128 // 4,
feed_forward_dimensions=256,
attention_type="full",
final_normalization=False).get()
t1.eval()
t2.eval()
with self.assertRaises(RuntimeError):
t2.load_state_dict(t1.state_dict())
t2.load_state_dict(PytorchMapper().map(t1.state_dict()))
x = torch.rand(3, 10, 128)
o1 = t2(x)
o2 = t1(x.permute(1, 0, 2)).permute(1, 0, 2)
self.assertLess(torch.abs(o1 - o2).max().item(), 1e-5)
def __init__(self,
d_model,
coords, flatten_order,
attention_type="linear",
n_layers=4,
n_heads=4,
d_query=32,
dropout=0.1,
attention_dropout=0.1):
super(SResTransformerTrain, self).__init__()
self.fourier_coefficient_embedding = torch.nn.Linear(2, d_model // 2)
self.pos_embedding = PositionalEncoding2D(
d_model // 2,
coords=coords,
flatten_order=flatten_order,
persistent=False
)
self.encoder = TransformerEncoderBuilder.from_kwargs(
attention_type=attention_type,
n_layers=n_layers,
n_heads=n_heads,
feed_forward_dimensions=n_heads * d_query * 4,
query_dimensions=d_query,
value_dimensions=d_query,
dropout=dropout,
attention_dropout=attention_dropout
).get()
self.predictor_amp = torch.nn.Linear(
n_heads * d_query,
1
)
self.predictor_phase = torch.nn.Linear(
n_heads * d_query,
1
)
def __init__(self,
d_model,
sequence_length,
mixtures,
attention_type="full",
n_layers=4,
n_heads=4,
d_query=32,
dropout=0.1,
softmax_temp=None,
attention_dropout=0.1,
bits=32,
rounds=4,
chunk_size=32,
masked=True):
super(ImageGenerator, self).__init__()
self.pos_embedding = PositionalEncoding(d_model // 2,
max_len=sequence_length)
self.value_embedding = torch.nn.Embedding(256, d_model // 2)
self.transformer = TransformerEncoderBuilder.from_kwargs(
attention_type=attention_type,
n_layers=n_layers,
n_heads=n_heads,
feed_forward_dimensions=n_heads * d_query * 4,
query_dimensions=d_query,
value_dimensions=d_query,
dropout=dropout,
softmax_temp=softmax_temp,
attention_dropout=attention_dropout,
bits=bits,
rounds=rounds,
chunk_size=chunk_size,
masked=masked).get()
hidden_size = n_heads * d_query
self.predictor = torch.nn.Linear(hidden_size, mixtures * 3)
def test_qkv(self):
d = {}
def store_qkv(event):
d["q"] = event.queries
d["k"] = event.keys
d["v"] = event.values
# default transformer is 4 layers 4 heads
transformer = TransformerEncoderBuilder().get()
x = transformer(torch.rand(1, 100, 64*4))
self.assertEqual(len(d), 0)
EventDispatcher.get().listen(QKVEvent, store_qkv)
x = transformer(torch.rand(1, 100, 64*4))
self.assertEqual(len(d), 3)
d.clear()
EventDispatcher.get().remove(store_qkv)
x = transformer(torch.rand(1, 100, 64*4))
self.assertEqual(len(d), 0)
d.clear()
EventDispatcher.get().listen(
QKVEvent & layer_name_contains(transformer, "layers.2.attention"),
store_qkv
)
x = transformer(torch.rand(1, 100, 64*4))
self.assertEqual(len(d), 3)
d.clear()
EventDispatcher.get().listen(
QKVEvent & layer_name_contains(transformer, "layers.22.attention"),
store_qkv
)
x = transformer(torch.rand(1, 100, 64*4))
self.assertEqual(len(d), 0)
d.clear()
EventDispatcher.get().clear()
def __init__(self,
input_shape,
bins,
width=128,
depth=2,
heads=1,
attn_dropout=0.0,
resid_dropout=0.0,
emb_dropout=0.0,
mask=True,
zero_out=False,
pos_init=False,
x_cond=False,
y_cond=False,
encoder_dims=0,
only_encode=False,
merged_decoder=False,
prime_len=None,
m_attn=0.25,
m_mlp=1,
init_scale=1.0,
checkpoint_res=0,
train=False):
super().__init__()
self.input_shape = input_shape
self.input_dims = input_dims = np.prod(input_shape)
self.encoder_dims = encoder_dims
self.bins = bins
self.width = width
self.depth = depth
self.x_emb = nn.Embedding(bins, width)
nn.init.normal_(self.x_emb.weight, std=0.02 * init_scale)
self.x_emb_dropout = nn.Dropout(emb_dropout)
self.y_cond = y_cond
self.x_cond = x_cond
if not y_cond:
self.start_token = nn.Parameter(
get_normal(1, width, std=0.01 * init_scale))
self.pos_emb = PositionEmbedding(input_shape=input_shape,
width=width,
init_scale=init_scale,
pos_init=pos_init)
self.pos_emb_dropout = nn.Dropout(emb_dropout)
if train:
self.transformer = TransformerEncoderBuilder.from_kwargs(
n_layers=depth,
n_heads=heads,
feed_forward_dimensions=int(m_mlp * width),
model_dimensions=width,
query_dimensions=int(m_attn * width) // heads,
value_dimensions=int(m_attn * width) // heads,
activation='gelu',
dropout=attn_dropout,
attention_type="causal-linear",
).get()
else:
# encoder (inference)
print(' [o] using RNN backend.')
self.transformer = RecurrentEncoderBuilder.from_kwargs(
n_layers=depth,
n_heads=heads,
model_dimensions=width,
feed_forward_dimensions=int(m_mlp * width),
query_dimensions=int(m_attn * width) // heads,
value_dimensions=int(m_attn * width) // heads,
dropout=attn_dropout,
activation='gelu',
attention_type="causal-linear",
).get()
self.only_encode = only_encode
self.prime_len = prime_len
if merged_decoder:
# Merged piped model uses this setup
self.add_cond_after_transformer = False
self.share_x_emb_x_out = False
else:
self.add_cond_after_transformer = True
self.share_x_emb_x_out = True
if not only_encode:
self.x_out = nn.Linear(width, bins, bias=False)
if self.share_x_emb_x_out:
self.x_out.weight = self.x_emb.weight
self.loss = t.nn.CrossEntropyLoss()
def test_attention_parameter(self):
builder = TransformerEncoderBuilder()
builder.n_layers = 3
builder.n_heads = 4
builder.feed_forward_dimensions = 512
builder.query_dimensions = 32
builder.value_dimensions = 64
builder.dropout = 0.1
builder.activation = "relu"
builder.final_normalization = True
# Full attention parameters
builder.softmax_temp = 1.0
builder.attention_dropout = 0.1
# Linear attention parameters
builder.feature_map = lambda x: (x > 0).float() * x
# Clustered attention parameters
builder.clusters = 100
builder.iterations = 10
builder.bits = 32
builder.hash_bias = True
# Exact topk attention parameters
builder.topk = 32
# Conditional attention parameters
builder.length_limit = 512
# Reformer attention parameters
builder.chunk_size = 32
builder.rounds = 1
# Add here old parameters to avoid regressions
invalid = ["dropout_rate"]
for name in invalid:
with self.assertRaises(AttributeError):
setattr(builder, name, None)
def __init__(self, n_token, is_training=True):
super(TransformerModel, self).__init__()
# --- params config --- #
self.n_token = n_token
self.d_model = D_MODEL
self.n_layer = N_LAYER #
self.dropout = 0.1
self.n_head = N_HEAD #
self.d_head = D_MODEL // N_HEAD
self.d_inner = 2048
self.loss_func = nn.CrossEntropyLoss(reduction='none')
self.emb_sizes = [128, 256, 64, 32, 512, 128, 128]
# --- modules config --- #
# embeddings
print('>>>>>:', self.n_token)
self.word_emb_tempo = Embeddings(self.n_token[0], self.emb_sizes[0])
self.word_emb_chord = Embeddings(self.n_token[1], self.emb_sizes[1])
self.word_emb_barbeat = Embeddings(self.n_token[2], self.emb_sizes[2])
self.word_emb_type = Embeddings(self.n_token[3], self.emb_sizes[3])
self.word_emb_pitch = Embeddings(self.n_token[4], self.emb_sizes[4])
self.word_emb_duration = Embeddings(self.n_token[5], self.emb_sizes[5])
self.word_emb_velocity = Embeddings(self.n_token[6], self.emb_sizes[6])
self.pos_emb = PositionalEncoding(self.d_model, self.dropout)
# linear
self.in_linear = nn.Linear(np.sum(self.emb_sizes), self.d_model)
# encoder
if is_training:
# encoder (training)
self.transformer_encoder = TransformerEncoderBuilder.from_kwargs(
n_layers=self.n_layer,
n_heads=self.n_head,
query_dimensions=self.d_model // self.n_head,
value_dimensions=self.d_model // self.n_head,
feed_forward_dimensions=2048,
activation='gelu',
dropout=0.1,
attention_type="causal-linear",
).get()
else:
# encoder (inference)
print(' [o] using RNN backend.')
self.transformer_encoder = RecurrentEncoderBuilder.from_kwargs(
n_layers=self.n_layer,
n_heads=self.n_head,
query_dimensions=self.d_model // self.n_head,
value_dimensions=self.d_model // self.n_head,
feed_forward_dimensions=2048,
activation='gelu',
dropout=0.1,
attention_type="causal-linear",
).get()
# blend with type
self.project_concat_type = nn.Linear(self.d_model + 32, self.d_model)
# individual output
self.proj_tempo = nn.Linear(self.d_model, self.n_token[0])
self.proj_chord = nn.Linear(self.d_model, self.n_token[1])
self.proj_barbeat = nn.Linear(self.d_model, self.n_token[2])
self.proj_type = nn.Linear(self.d_model, self.n_token[3])
self.proj_pitch = nn.Linear(self.d_model, self.n_token[4])
self.proj_duration = nn.Linear(self.d_model, self.n_token[5])
self.proj_velocity = nn.Linear(self.d_model, self.n_token[6])
def __init__(self,
embedding_dim,
hidden_dim,
problem,
attention_type,
n_encode_layers=2,
feed_forward_dim=512,
tanh_clipping=10.,
mask_inner=True,
mask_logits=True,
normalization='batch',
n_heads=8,
encoding_knn_size=None,
decoding_knn_size=None,
checkpoint_encoder=False,
shrink_size=None):
super(AttentionModel, self).__init__()
self.embedding_dim = embedding_dim
self.hidden_dim = hidden_dim
self.n_encode_layers = n_encode_layers
self.decode_type = None
self.temp = 1.0
self.allow_partial = problem.NAME == 'sdvrp'
self.is_vrp = problem.NAME == 'cvrp' or problem.NAME == 'sdvrp'
self.is_orienteering = problem.NAME == 'op'
self.is_pctsp = problem.NAME == 'pctsp'
self.tanh_clipping = tanh_clipping
self.mask_inner = mask_inner
self.mask_logits = mask_logits
self.problem = problem
self.n_heads = n_heads
self.checkpoint_encoder = checkpoint_encoder
self.shrink_size = shrink_size
# Problem specific context parameters (placeholder and step context dimension)
if self.is_vrp or self.is_orienteering or self.is_pctsp:
# Embedding of last node + remaining_capacity / remaining length / remaining prize to collect
step_context_dim = embedding_dim + 1
if self.is_pctsp:
node_dim = 4 # x, y, expected_prize, penalty
else:
node_dim = 3 # x, y, demand / prize
# Special embedding projection for depot node
self.init_embed_depot = nn.Linear(2, embedding_dim)
if self.is_vrp and self.allow_partial: # Need to include the demand if split delivery allowed
self.project_node_step = nn.Linear(1,
3 * embedding_dim,
bias=False)
else: # TSP
assert problem.NAME == "tsp", "Unsupported problem: {}".format(
problem.NAME)
step_context_dim = 2 * embedding_dim # Embedding of first and last node
node_dim = 2 # x, y
# Learned input symbols for first action
self.W_placeholder = nn.Parameter(torch.Tensor(2 * embedding_dim))
self.W_placeholder.data.uniform_(
-1, 1) # Placeholder should be in range of activations
self.encoding_knn_size = encoding_knn_size
self.decoding_knn_size = decoding_knn_size
self.init_embed = nn.Linear(node_dim, embedding_dim)
self.attention_type = attention_type
if attention_type == 'original':
self.embedder = GraphAttentionEncoder(
n_heads=n_heads,
embed_dim=embedding_dim,
feed_forward_dim=feed_forward_dim,
n_layers=n_encode_layers,
normalization=normalization)
else:
self.embedder = TransformerEncoderBuilder.from_kwargs(
n_layers=n_encode_layers,
n_heads=n_heads,
query_dimensions=embedding_dim // n_heads,
value_dimensions=embedding_dim // n_heads,
feed_forward_dimensions=512,
attention_dropout=0.0,
local_context=20,
clusters=5,
topk=20,
feature_map=Favor.factory(n_dims=128),
attention_type=attention_type).get()
# For each node we compute (glimpse key, glimpse value, logit key) so 3 * embedding_dim
self.project_node_embeddings = nn.Linear(embedding_dim,
3 * embedding_dim,
bias=False)
self.project_fixed_context = nn.Linear(embedding_dim,
embedding_dim,
bias=False)
self.project_step_context = nn.Linear(step_context_dim,
embedding_dim,
bias=False)
assert embedding_dim % n_heads == 0
# Note n_heads * val_dim == embedding_dim so input to project_out is embedding_dim
self.project_out = nn.Linear(embedding_dim, embedding_dim, bias=False)
def __init__(self,
num_feats,
num_output_points,
lstm_layers,
n_layers,
n_heads,
hidden_dim,
ff_dim,
tf_depth=3,
dropout=0.15):
super(SetTransformer, self).__init__()
self.num_feats = num_feats
self.k = num_output_points
def dup(x):
return (x, x) if type(x) == int else x
self.lstm_layers = lstm_layers
self.n_layers = dup(n_layers)
self.n_heads = dup(n_heads)
self.hidden_dim = dup(hidden_dim)
self.ff_dim = dup(ff_dim)
self.tf_depth = dup(tf_depth)
self.d_model = [self.hidden_dim[i] * self.n_heads[i] for i in [0, 1]]
encoder_builder_pre = TransformerEncoderBuilder.from_kwargs(
n_layers=self.n_layers[0],
n_heads=self.n_heads[0],
query_dimensions=self.hidden_dim[0],
value_dimensions=self.hidden_dim[0],
feed_forward_dimensions=self.ff_dim[0],
attention_type='linear',
dropout=dropout)
encoder_builder_post = TransformerEncoderBuilder.from_kwargs(
n_layers=self.n_layers[1],
n_heads=self.n_heads[1],
query_dimensions=self.hidden_dim[1],
value_dimensions=self.hidden_dim[1],
feed_forward_dimensions=self.ff_dim[1],
attention_type='linear',
dropout=dropout)
self.seeds = nn.Parameter(torch.normal(0, 1,
(self.k, self.d_model[0])))
self.encoder_pre = encoder_builder_pre.get()
self.encoder_post = encoder_builder_post.get()
self.initial_ff = nn.Linear(self.num_feats, self.d_model[0])
# self.pos_encoding = PositionalEncoding(self.d_model[0], dropout=dropout)
self.lstm = nn.LSTM(self.d_model[0],
self.d_model[0],
2,
batch_first=True,
bidirectional=False)
self.attn_pooling = AttentionLayer(LinearAttention(self.d_model[0]),
self.d_model[0], self.n_heads[0])
self.final_ff = nn.Linear(self.d_model[1], self.num_feats)
# init masks to meaningless values, doesn't matter what. these are all empty anyway.
self.mask = FullMask(N=self.k, M=5)
self.kl_mask = LengthMask(torch.ones(5) * 5)
self.ql_mask = LengthMask(torch.ones(5) * self.k)
