Spaces:
Build error
Build error
File size: 3,312 Bytes
7eb6194 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 |
# -*- coding: utf-8 -*-
# @Author : Lintao Peng
# @File : SGFMT.py
# coding=utf-8
# Design based on the Vit
import torch.nn as nn
from net.IntmdSequential import IntermediateSequential
#实现了自注意力机制，相当于unet的bottleneck层
class SelfAttention(nn.Module):
def __init__(
self, dim, heads=8, qkv_bias=False, qk_scale=None, dropout_rate=0.0
):
super().__init__()
self.num_heads = heads
head_dim = dim // heads
self.scale = qk_scale or head_dim ** -0.5
self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
self.attn_drop = nn.Dropout(dropout_rate)
self.proj = nn.Linear(dim, dim)
self.proj_drop = nn.Dropout(dropout_rate)
def forward(self, x):
B, N, C = x.shape
qkv = (
self.qkv(x)
.reshape(B, N, 3, self.num_heads, C // self.num_heads)
.permute(2, 0, 3, 1, 4)
)
q, k, v = (
qkv[0],
qkv[1],
qkv[2],
) # make torchscript happy (cannot use tensor as tuple)
attn = (q @ k.transpose(-2, -1)) * self.scale
attn = attn.softmax(dim=-1)
attn = self.attn_drop(attn)
x = (attn @ v).transpose(1, 2).reshape(B, N, C)
x = self.proj(x)
x = self.proj_drop(x)
return x
class Residual(nn.Module):
def __init__(self, fn):
super().__init__()
self.fn = fn
def forward(self, x):
return self.fn(x) + x
class PreNorm(nn.Module):
def __init__(self, dim, fn):
super().__init__()
self.norm = nn.LayerNorm(dim)
self.fn = fn
def forward(self, x):
return self.fn(self.norm(x))
class PreNormDrop(nn.Module):
def __init__(self, dim, dropout_rate, fn):
super().__init__()
self.norm = nn.LayerNorm(dim)
self.dropout = nn.Dropout(p=dropout_rate)
self.fn = fn
def forward(self, x):
return self.dropout(self.fn(self.norm(x)))
class FeedForward(nn.Module):
def __init__(self, dim, hidden_dim, dropout_rate):
super().__init__()
self.net = nn.Sequential(
nn.Linear(dim, hidden_dim),
nn.GELU(),
nn.Dropout(p=dropout_rate),
nn.Linear(hidden_dim, dim),
nn.Dropout(p=dropout_rate),
)
def forward(self, x):
return self.net(x)
class TransformerModel(nn.Module):
def __init__(
self,
dim, #512
depth, #4
heads, #8
mlp_dim, #4096
dropout_rate=0.1,
attn_dropout_rate=0.1,
):
super().__init__()
layers = []
for _ in range(depth):
layers.extend(
[
Residual(
PreNormDrop(
dim,
dropout_rate,
SelfAttention(dim, heads=heads, dropout_rate=attn_dropout_rate),
)
),
Residual(
PreNorm(dim, FeedForward(dim, mlp_dim, dropout_rate))
),
]
)
# dim = dim / 2
self.net = IntermediateSequential(*layers)
def forward(self, x):
return self.net(x)
|