146 lines
5.1 KiB
Python
146 lines
5.1 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
class StochDepth(nn.Module):
|
|
"""Batchwise Dropout used in EfficientNet, optionally sans rescaling."""
|
|
|
|
def __init__(self, drop_rate, scale_by_keep=False):
|
|
super().__init__()
|
|
self.drop_rate = drop_rate
|
|
self.scale_by_keep = scale_by_keep
|
|
|
|
def forward(self, x):
|
|
if not self.training:
|
|
return x
|
|
|
|
batch_size = x.shape[0]
|
|
r = torch.rand([batch_size, 1, 1], dtype=x.dtype, device=x.device)
|
|
keep_prob = 1.0 - self.drop_rate
|
|
binary_tensor = torch.floor(keep_prob + r)
|
|
if self.scale_by_keep:
|
|
x = x / keep_prob
|
|
return x * binary_tensor
|
|
|
|
class PosEmbed(nn.Module):
|
|
def __init__(self, input_size):
|
|
super().__init__()
|
|
|
|
self.pos_embed = nn.Parameter(
|
|
torch.empty(input_size, dtype=torch.float32)
|
|
)
|
|
torch.nn.init.trunc_normal_(self.pos_embed,mean=0.0, std=0.02)
|
|
|
|
def forward(self, x):
|
|
return x + self.pos_embed.unsqueeze(0)
|
|
|
|
class MLPBlock(nn.Module):
|
|
def __init__(self, input_dim, mlp_dim, stochdepth_rate):
|
|
super().__init__()
|
|
self.stochdepth_rate = stochdepth_rate
|
|
self.fc1 = nn.Linear(input_dim, mlp_dim)
|
|
self.fc2 = nn.Linear(mlp_dim, input_dim)
|
|
if stochdepth_rate > 0.0:
|
|
self.stochdepth = StochDepth(stochdepth_rate, scale_by_keep=True)
|
|
else:
|
|
self.stochdepth = None
|
|
|
|
def forward(self, x):
|
|
out = F.gelu(self.fc1(x))
|
|
if self.stochdepth:
|
|
out = self.stochdepth(out)
|
|
out = self.fc2(out)
|
|
return out
|
|
|
|
class SkipInitChannelwise(nn.Module):
|
|
def __init__(self, channels, init_val=1e-6):
|
|
super().__init__()
|
|
self.channels = channels
|
|
self.init_val = init_val
|
|
self.skip = nn.Parameter(torch.ones(channels) * init_val)
|
|
|
|
def forward(self, x):
|
|
return x * self.skip
|
|
|
|
class ViTBlock(nn.Module):
|
|
def __init__(self, input_dim, heads, key_dim, mlp_dim, layerscale_init, stochdepth_rate):
|
|
super().__init__()
|
|
self.norm1 = nn.LayerNorm(input_dim,eps=1e-3)
|
|
self.attn = nn.MultiheadAttention(input_dim, heads, batch_first=True)
|
|
self.skip1 = SkipInitChannelwise(key_dim, init_val=layerscale_init)
|
|
self.stochdepth1 = StochDepth(stochdepth_rate, scale_by_keep=True) if stochdepth_rate > 0.0 else None
|
|
self.norm2 = nn.LayerNorm(input_dim,eps=1e-3)
|
|
self.mlp = MLPBlock(input_dim, mlp_dim, stochdepth_rate)
|
|
self.skip2 = SkipInitChannelwise(key_dim, init_val=layerscale_init)
|
|
self.stochdepth2 = StochDepth(stochdepth_rate, scale_by_keep=True) if stochdepth_rate > 0.0 else None
|
|
|
|
def forward(self, x):
|
|
out = self.norm1(x)
|
|
out = self.attn(out, out, out)[0]
|
|
out = self.skip1(out)
|
|
if self.stochdepth1:
|
|
out = self.stochdepth1(out)
|
|
x = out + x
|
|
|
|
out = self.norm2(x)
|
|
out = self.mlp(out)
|
|
out = self.skip2(out)
|
|
if self.stochdepth2:
|
|
out = self.stochdepth2(out)
|
|
|
|
out = out + x
|
|
return out
|
|
|
|
class ViT(nn.Module):
|
|
def __init__(self, in_channels=3, img_size=320, out_classes=2000, definition_name="B16"):
|
|
super().__init__()
|
|
self.definitions = {
|
|
"B16": {
|
|
"num_blocks": 12,
|
|
"patch_size": 16,
|
|
"key_dim": 768,
|
|
"mlp_dim": 3072,
|
|
"heads": 12,
|
|
"stochdepth_rate": 0.05,
|
|
},
|
|
# Other definitions removed for simplicity
|
|
}
|
|
|
|
definition = self.definitions[definition_name]
|
|
self.blocks = nn.ModuleList()
|
|
num_blocks = definition["num_blocks"]
|
|
patch_size = definition["patch_size"]
|
|
key_dim = definition["key_dim"]
|
|
mlp_dim = definition["mlp_dim"]
|
|
heads = definition["heads"]
|
|
stochdepth_rate = definition["stochdepth_rate"]
|
|
layerscale_init = 0.1 # Replacing CaiT_LayerScale_init(num_blocks)
|
|
|
|
self.conv = nn.Conv2d(in_channels, key_dim, kernel_size=patch_size, stride=patch_size)
|
|
self.pos_embed = PosEmbed(((img_size // patch_size) ** 2,key_dim))
|
|
|
|
for i in range(num_blocks):
|
|
self.blocks.append(
|
|
ViTBlock(key_dim, heads, key_dim, mlp_dim, layerscale_init, stochdepth_rate)
|
|
)
|
|
|
|
self.norm = nn.LayerNorm(key_dim,eps=1e-3)
|
|
self.avgpool = nn.AdaptiveAvgPool1d(1)
|
|
self.fc = nn.Linear(key_dim, out_classes)
|
|
self.act = nn.Sigmoid()
|
|
|
|
def forward(self, x):
|
|
x = (x - 127.5) / 127.5
|
|
x = self.conv(x)
|
|
b, c, h, w = x.shape
|
|
x = x.view(b, c, h*w).permute(0, 2, 1) # (B, H*W, C)
|
|
x = self.pos_embed(x)
|
|
for block in self.blocks:
|
|
x = block(x)
|
|
x = self.norm(x)
|
|
x = self.avgpool(x.transpose(1, 2)).squeeze(-1) # (B, C)
|
|
|
|
# pfg uses output of last pooling layer.
|
|
#x = self.fc(x)
|
|
#x = self.act(x)
|
|
return x |