以下是一个符合要求的VIT五分类网络：
```python
import torch
import torch.nn as nn
import torch.nn.functional as F
class MLP(nn.Module):
def __init__(self, in_features, hidden_features, out_features):
super().__init__()
self.fc1 = nn.Linear(in_features, hidden_features)
self.fc2 = nn.Linear(hidden_features, out_features)
def forward(self, x):
x = F.relu(self.fc1(x))
x = F.dropout(x, p=0.1, training=self.training)
x = self.fc2(x)
return x
class Attention(nn.Module):
def __init__(self, dim, heads=8, dropout=0.1):
super().__init__()
self.heads = heads
self.scale = dim ** -0.5
self.to_qkv = nn.Linear(dim, dim * 3, bias=False)
self.to_out = nn.Linear(dim, dim)
self.dropout = nn.Dropout(dropout)
def forward(self, x):
b, n, _, h = *x.shape, self.heads
qkv = self.to_qkv(x).chunk(3, dim=-1)
q, k, v = map(lambda t: t.reshape(b, n, h, -1).transpose(1, 2), qkv)
dots = torch.einsum('bhid,bhjd->bhij', q, k) * self.scale
attn = dots.softmax(dim=-1)
attn = self.dropout(attn)
out = torch.einsum('bhij,bhjd->bhid', attn, v)
out = out.transpose(1, 2).reshape(b, n, -1)
out = self.to_out(out)
out = self.dropout(out)
return out
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 FeedForward(nn.Module):
def __init__(self, dim, hidden_dim, dropout=0.1):
super().__init__()
self.net = nn.Sequential(
nn.Linear(dim, hidden_dim),
nn.GELU(),
nn.Dropout(dropout),
nn.Linear(hidden_dim, dim),
nn.Dropout(dropout)
)
def forward(self, x):
return self.net(x)
class CBAM(nn.Module):
def __init__(self, in_features, reduction_ratio=16):
super().__init__()
self.in_features = in_features
self.reduction_ratio = reduction_ratio
self.avg_pool = nn.AdaptiveAvgPool2d((1,1))
self.max_pool = nn.AdaptiveMaxPool2d((1,1))
self.fc1 = nn.Linear(in_features, in_features // reduction_ratio)
self.relu = nn.ReLU()
self.fc2 = nn.Linear(in_features // reduction_ratio, in_features)
self.sigmoid = nn.Sigmoid()
def forward(self, x):
b, c, _, _ = x.size()
avg_out = self.avg_pool(x).view(b, c)
avg_out = self.fc2(self.relu(self.fc1(avg_out)))
max_out = self.max_pool(x).view(b, c)
max_out = self.fc2(self.relu(self.fc1(max_out)))
out = avg_out + max_out
out = self.sigmoid(out).view(b, c, 1, 1)
return x * out
class VIT(nn.Module):
def __init__(self, *, image_size, patch_size, num_classes, dim, depth, heads, mlp_dim, dropout):
super().__init__()
assert image_size % patch_size == 0, 'image size must be divisible by patch size'
num_patches = (image_size // patch_size) ** 2
patch_dim = 3 * patch_size ** 2
self.patch_size = patch_size
self.pos_embedding = nn.Parameter(torch.randn(1, num_patches + 1, dim))
self.patch_to_embedding = nn.Linear(patch_dim, dim)
self.cls_token = nn.Parameter(torch.randn(1, 1, dim))
self.dropout = nn.Dropout(dropout)
self.transformer = nn.ModuleList([
Residual(PreNorm(dim, Attention(dim, heads=heads, dropout=dropout))),
Residual(PreNorm(dim, FeedForward(dim, mlp_dim, dropout=dropout)))
])
self.to_cls_token = nn.Identity()
self.mlp_head = MLP(dim, hidden_features=512, out_features=num_classes)
self.cbam = CBAM(in_features=dim)
def forward(self, x):
b, _, h, w = x.shape
p = self.patch_size
# assert input size is divisible by patch size
assert h % p == 0 and w % p == 0, f'image size {h}x{w} not divisible by patch size {p}'
# convert image to patches
x = x.reshape(b, 3, h//p, p, w//p, p)
x = x.permute(0, 2, 4, 1, 3, 5)
x = x.reshape(b, -1, 3 * p ** 2)
# add cls token and convert patches to embeddings
cls_token = self.cls_token.expand(b, -1, -1)
x = torch.cat((cls_token, x), dim=1)
x = self.patch_to_embedding(x)
# add positional embedding
x = x + self.pos_embedding
# apply transformer
for transformer_block in self.transformer:
x = transformer_block(x)
# extract cls token and apply MLP head
x = self.to_cls_token(x[:, 0])
x = self.cbam(x.permute(0, 2, 1).reshape(b, -1, h//p, w//p)).permute(0, 2, 3, 1).reshape(b, -1)
x = self.dropout(x)
x = self.mlp_head(x)
return x
```
使用方式：
```python
import torch
model = VIT(image_size=128, patch_size=16, num_classes=5, dim=256, depth=6, heads=8, mlp_dim=512, dropout=0.1)
input_tensor = torch.randn(64, 3, 128, 128)
output_tensor = model(input_tensor)
print(output_tensor.shape) # 输出 (64, 5)
```
最低0.47元/天 解锁文章
876
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
