Spaces:
Running
Running
File size: 4,395 Bytes
c608946 9cde3b4 c608946 2947428 c608946 2947428 c608946 2947428 c608946 9cde3b4 c608946 2947428 c608946 2947428 c608946 2947428 c608946 2947428 c608946 9cde3b4 c608946 9cde3b4 c608946 2947428 c608946 2947428 c608946 9cde3b4 c608946 2947428 c608946 2947428 c608946 9cde3b4 c608946 2947428 c608946 2947428 c608946 2947428 c608946 2947428 c608946 2947428 |
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 |
from typing import Optional, Union
import torch
from torch import device
import torch.nn as nn
import torch.nn.functional as F
import torchvision.models as tvm
import gc
class ResNet50(nn.Module):
def __init__(self, pretrained=False, high_res = False, weights = None,
dilation = None, freeze_bn = True, anti_aliased = False, early_exit = False, amp = False, amp_dtype = torch.float16) -> None:
super().__init__()
if dilation is None:
dilation = [False,False,False]
if anti_aliased:
pass
else:
if weights is not None:
self.net = tvm.resnet50(weights = weights,replace_stride_with_dilation=dilation)
else:
self.net = tvm.resnet50(pretrained=pretrained,replace_stride_with_dilation=dilation)
self.high_res = high_res
self.freeze_bn = freeze_bn
self.early_exit = early_exit
self.amp = amp
self.amp_dtype = amp_dtype
def forward(self, x, **kwargs):
with torch.autocast("cuda", enabled=self.amp, dtype = self.amp_dtype):
net = self.net
feats = {1:x}
x = net.conv1(x)
x = net.bn1(x)
x = net.relu(x)
feats[2] = x
x = net.maxpool(x)
x = net.layer1(x)
feats[4] = x
x = net.layer2(x)
feats[8] = x
if self.early_exit:
return feats
x = net.layer3(x)
feats[16] = x
x = net.layer4(x)
feats[32] = x
return feats
def train(self, mode=True):
super().train(mode)
if self.freeze_bn:
for m in self.modules():
if isinstance(m, nn.BatchNorm2d):
m.eval()
pass
class VGG19(nn.Module):
def __init__(self, pretrained=False, amp = False, amp_dtype = torch.float16) -> None:
super().__init__()
self.layers = nn.ModuleList(tvm.vgg19_bn(pretrained=pretrained).features[:40])
self.amp = amp
self.amp_dtype = amp_dtype
def forward(self, x, **kwargs):
with torch.autocast("cuda", enabled=self.amp, dtype = self.amp_dtype):
feats = {}
scale = 1
for layer in self.layers:
if isinstance(layer, nn.MaxPool2d):
feats[scale] = x
scale = scale*2
x = layer(x)
return feats
class CNNandDinov2(nn.Module):
def __init__(self, cnn_kwargs = None, amp = False, use_vgg = False, dinov2_weights = None, amp_dtype = torch.float16):
super().__init__()
if dinov2_weights is None:
dinov2_weights = torch.hub.load_state_dict_from_url("https://dl.fbaipublicfiles.com/dinov2/dinov2_vitl14/dinov2_vitl14_pretrain.pth", map_location="cpu")
from .transformer import vit_large
vit_kwargs = dict(img_size= 518,
patch_size= 14,
init_values = 1.0,
ffn_layer = "mlp",
block_chunks = 0,
)
dinov2_vitl14 = vit_large(**vit_kwargs).eval()
dinov2_vitl14.load_state_dict(dinov2_weights)
cnn_kwargs = cnn_kwargs if cnn_kwargs is not None else {}
if not use_vgg:
self.cnn = ResNet50(**cnn_kwargs)
else:
self.cnn = VGG19(**cnn_kwargs)
self.amp = amp
self.amp_dtype = amp_dtype
if self.amp:
dinov2_vitl14 = dinov2_vitl14.to(self.amp_dtype)
self.dinov2_vitl14 = [dinov2_vitl14] # ugly hack to not show parameters to DDP
def train(self, mode: bool = True):
return self.cnn.train(mode)
def forward(self, x, upsample = False):
B,C,H,W = x.shape
feature_pyramid = self.cnn(x)
if not upsample:
with torch.no_grad():
if self.dinov2_vitl14[0].device != x.device:
self.dinov2_vitl14[0] = self.dinov2_vitl14[0].to(x.device).to(self.amp_dtype)
dinov2_features_16 = self.dinov2_vitl14[0].forward_features(x.to(self.amp_dtype))
features_16 = dinov2_features_16['x_norm_patchtokens'].permute(0,2,1).reshape(B,1024,H//14, W//14)
del dinov2_features_16
feature_pyramid[16] = features_16
return feature_pyramid |