def __init__(self, cfg):
super(PCNN, self).__init__()
self.use_pcnn = cfg.use_pcnn
if cfg.dim_strategy == 'cat':
cfg.in_channels = cfg.word_dim + 2 * cfg.pos_dim
else:
cfg.in_channels = cfg.word_dim
self.embedding = Embedding(cfg)
self.cnn = CNN(cfg)
self.fc1 = nn.Linear(
len(cfg.kernel_sizes) * cfg.out_channels, cfg.intermediate)
self.fc2 = nn.Linear(cfg.intermediate, cfg.num_relations)
self.dropout = nn.Dropout(cfg.dropout)
if self.use_pcnn:
self.fc_pcnn = nn.Linear(
3 * len(cfg.kernel_sizes) * cfg.out_channels,
len(cfg.kernel_sizes) * cfg.out_channels)
self.pcnn_mask_embedding = nn.Embedding(4, 3)
masks = torch.tensor([[0, 0, 0], [100, 0, 0], [0, 100, 0],
[0, 0, 100]])
self.pcnn_mask_embedding.weight.data.copy_(masks)
self.pcnn_mask_embedding.weight.requires_grad = False
def test_CNN():
x = torch.randn(4, 5, 100)
seq = torch.arange(4, 0, -1)
mask = seq_len_to_mask(seq, max_len=5)
cnn = CNN(config)
out, out_pooling = cnn(x, mask=mask)
out_channels = config.out_channels * len(config.kernel_sizes)
assert out.shape == torch.Size([4, 5, out_channels])
assert out_pooling.shape == torch.Size([4, out_channels])
def __init__(self, cfg):
super(Capsule, self).__init__()
if cfg.dim_strategy == 'cat':
cfg.in_channels = cfg.word_dim + 2 * cfg.pos_dim
else:
cfg.in_channels = cfg.word_dim
# capsule config
cfg.input_dim_capsule = cfg.out_channels
cfg.num_capsule = cfg.num_relations
self.num_relations = cfg.num_relations
self.embedding = Embedding(cfg)
self.cnn = CNN(cfg)
self.capsule = CapsuleLayer(cfg)