| import torch |
| from torch import nn |
|
|
|
|
| class ValueEmbedding(nn.Module): |
| def __init__(self, hidden_dim, do_fft): |
| super(ValueEmbedding, self).__init__() |
| self.embedding = nn.Linear(1, hidden_dim) |
| self.do_fft = do_fft |
|
|
| def forward(self, x): |
| if self.do_fft: |
| x = torch.fft.fft(x, dim=-1).real |
| return self.embedding(x.unsqueeze(-1)) |
|
|
|
|
| class PositionalEmbeddingWithDnaPosition(nn.Module): |
| def __init__(self, d_model, batch_size, seq_length, positional_temp, device): |
| super(PositionalEmbeddingWithDnaPosition, self).__init__() |
| self.d_model = d_model |
| self.encoding = torch.zeros(batch_size, seq_length, self.d_model, device=device) |
| |
| self.encoding.requires_grad_(False) |
| self.positional_temp = positional_temp |
|
|
| def forward(self, x): |
| pos = x.float().unsqueeze(-1) / 1 |
| _2i = torch.arange(0, self.d_model, 2, device=x.device) |
| self.encoding[:x.shape[0], :, 0::2] = torch.sin(pos / self.positional_temp / (10000 ** (_2i / self.d_model))) |
| self.encoding[:x.shape[0], :, 1::2] = torch.cos(pos / self.positional_temp / (10000 ** (_2i / self.d_model))) |
| return self.encoding[:x.shape[0], :, :] |
|
|
|
|
| class PretrainEmbeddingSimple(nn.Module): |
| def __init__( |
| self, |
| embedding_dim, |
| chromosome_size, |
| embedding_dropout, |
| positional_embedding_type, |
| positional_temp, |
| batch_size, |
| seq_length, |
| device, |
| chromatin_embedding |
| ): |
| super(PretrainEmbeddingSimple, self).__init__() |
| self.value_embedding = ValueEmbedding(embedding_dim, False) |
| self.chromatin_embedding = chromatin_embedding |
| self.positional_embedding_type = positional_embedding_type |
| if self.chromatin_embedding: |
| self.chromosome_embedding = nn.Embedding(chromosome_size, embedding_dim) |
| self.position_embedding = PositionalEmbeddingWithDnaPosition( |
| embedding_dim, |
| batch_size, |
| seq_length, |
| positional_temp, |
| device |
| ) |
| self.embedding_dropout = embedding_dropout |
| self.dropout = nn.Dropout(p=self.embedding_dropout) |
|
|
| def forward(self, value, chromosome, hg38_start, hg38_end): |
| if self.chromatin_embedding: |
| embedding = self.value_embedding(value) + self.chromosome_embedding(chromosome) + self.position_embedding( |
| hg38_start) + self.position_embedding(hg38_end) |
| else: |
| embedding = self.value_embedding(value) |
| if self.embedding_dropout > 0: |
| return self.dropout(embedding) |
| else: |
| return embedding |
|
|
