以下是一个带有sparsity参数的稀疏Transformer的完整代码：
```python
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable
class SparseSelfAttention(nn.Module):
def __init__(self, d_model, num_heads, sparsity, dropout=0.1):
super(SparseSelfAttention, self).__init__()
self.d_model = d_model
self.num_heads = num_heads
self.sparsity = sparsity
self.dropout = nn.Dropout(dropout)
assert d_model % num_heads == 0
self.depth = d_model // num_heads
self.qkv = nn.Linear(d_model, 3 * d_model)
self.proj = nn.Linear(d_model, d_model)
self.softmax = nn.Softmax(dim=-1)
def forward(self, x, mask=None):
batch_size, seq_len, d_model = x.size()
qkv = self.qkv(x).view(batch_size, seq_len, 3, self.num_heads, self.depth)
qkv = qkv.permute(2, 0, 3, 1, 4)
q, k, v = qkv[0], qkv[1], qkv[2]
scores = torch.matmul(q, k.transpose(-2, -1))
scores = scores / (self.depth ** 0.5)
if mask is not None:
scores = scores.masked_fill(mask == 0, -1e9)
if self.sparsity > 0:
# Calculate sparsity pattern
num_sparse_tokens = int(seq_len * self.sparsity)
sparse_indices = torch.topk(scores.view(batch_size, self.num_heads, -1), num_sparse_tokens, dim=-1, largest=False)[1]
# Create sparse mask
mask = torch.ones_like(scores)
mask.scatter_(-1, sparse_indices, 0)
# Set sparse scores to -inf to exclude them from softmax
scores = scores.masked_fill(mask == 0, -1e9)
attn = self.softmax(scores)
attn = self.dropout(attn)
context = torch.matmul(attn, v)
context = context.permute(1, 2, 0, 3).contiguous().view(batch_size, seq_len, -1)
output = self.proj(context)
return output
class SparseTransformer(nn.Module):
def __init__(self, d_model, num_heads, num_layers, sparsity, dropout=0.1):
super(SparseTransformer, self).__init__()
self.d_model = d_model
self.num_heads = num_heads
self.num_layers = num_layers
self.sparsity = sparsity
self.dropout = nn.Dropout(dropout)
self.layers = nn.ModuleList([SparseSelfAttention(d_model, num_heads, sparsity, dropout) for _ in range(num_layers)])
def forward(self, x, mask=None):
for layer in self.layers:
x = x + layer(x, mask=mask)
x = self.dropout(x)
return x
```
上述代码中，在SparseSelfAttention类的初始化函数中增加了一个sparsity参数，用于控制稀疏程度。在forward函数中，当sparsity大于0时，会根据scores计算出一个稀疏模式，然后创建一个稀疏掩码，并将稀疏位置的scores设置为负无穷以排除它们的影响。通过这种方式，可以实现带有稀疏控制的Transformer模型。
本文探讨了网络中输入稀疏问题，通过介绍默认值设置、创新的稀疏卷积层及深度完成技术（如非引导深度上采样和图像引导的深度补全），以提升网络对低密度像素的有效利用。同时，关注了深度完成领域的两个关键点：一个关注像素有效性信息流，另一个聚焦于深度细节的恢复策略。
扫一扫
1734
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
