Spaces:
Runtime error
Runtime error
File size: 6,327 Bytes
8fc2b4e |
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 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 |
# Copyright (c) Aishwarya Kamath & Nicolas Carion. Licensed under the Apache License 2.0. All Rights Reserved
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
"""
Backbone modules.
"""
from collections import OrderedDict
import torch
import torch.nn.functional as F
import torchvision
from timm.models import create_model
from torch import nn
from torchvision.models._utils import IntermediateLayerGetter
from cliport.models.misc import NestedTensor
class FrozenBatchNorm2d(torch.nn.Module):
"""
BatchNorm2d where the batch statistics and the affine parameters are fixed.
Copy-paste from torchvision.misc.ops with added eps before rqsrt,
without which any other models than torchvision.models.resnet[18,34,50,101]
produce nans.
"""
def __init__(self, n):
super(FrozenBatchNorm2d, self).__init__()
self.register_buffer("weight", torch.ones(n))
self.register_buffer("bias", torch.zeros(n))
self.register_buffer("running_mean", torch.zeros(n))
self.register_buffer("running_var", torch.ones(n))
def _load_from_state_dict(
self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs
):
num_batches_tracked_key = prefix + "num_batches_tracked"
if num_batches_tracked_key in state_dict:
del state_dict[num_batches_tracked_key]
super(FrozenBatchNorm2d, self)._load_from_state_dict(
state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs
)
def forward(self, x):
# move reshapes to the beginning
# to make it fuser-friendly
w = self.weight.reshape(1, -1, 1, 1)
b = self.bias.reshape(1, -1, 1, 1)
rv = self.running_var.reshape(1, -1, 1, 1)
rm = self.running_mean.reshape(1, -1, 1, 1)
eps = 1e-5
scale = w * (rv + eps).rsqrt()
bias = b - rm * scale
return x * scale + bias
class BackboneBase(nn.Module):
def __init__(self, backbone: nn.Module, train_backbone: bool, num_channels: int, return_interm_layers: bool):
super().__init__()
for name, parameter in backbone.named_parameters():
if not train_backbone or "layer2" not in name and "layer3" not in name and "layer4" not in name:
parameter.requires_grad_(False)
if return_interm_layers:
return_layers = {"layer1": "0", "layer2": "1", "layer3": "2", "layer4": "3"}
else:
return_layers = {"layer4": 0}
self.body = IntermediateLayerGetter(backbone, return_layers=return_layers)
self.num_channels = num_channels
def forward(self, tensor_list):
xs = self.body(tensor_list.tensors)
out = OrderedDict()
for name, x in xs.items():
mask = F.interpolate(tensor_list.mask[None].float(), size=x.shape[-2:]).bool()[0]
out[name] = NestedTensor(x, mask)
return out
class Backbone(BackboneBase):
"""ResNet backbone with frozen BatchNorm."""
def __init__(self, name: str, train_backbone: bool, return_interm_layers: bool, dilation: bool):
backbone = getattr(torchvision.models, name)(
replace_stride_with_dilation=[False, False, dilation], pretrained=False, norm_layer=FrozenBatchNorm2d
)
num_channels = 512 if name in ("resnet18", "resnet34") else 2048
super().__init__(backbone, train_backbone, num_channels, return_interm_layers)
class GroupNorm32(torch.nn.GroupNorm):
def __init__(self, num_channels, num_groups=32, **kargs):
super().__init__(num_groups, num_channels, **kargs)
class GroupNormBackbone(BackboneBase):
"""ResNet backbone with GroupNorm with 32 channels."""
def __init__(self, name: str, train_backbone: bool, return_interm_layers: bool, dilation: bool):
name_map = {
"resnet50-gn": ("resnet50", "/checkpoint/szagoruyko/imagenet/22014122/checkpoint.pth"),
"resnet101-gn": ("resnet101", "/checkpoint/szagoruyko/imagenet/22080524/checkpoint.pth"),
}
backbone = getattr(torchvision.models, name_map[name][0])(
replace_stride_with_dilation=[False, False, dilation], pretrained=False, norm_layer=GroupNorm32
)
checkpoint = torch.load(name_map[name][1], map_location="cpu")
state_dict = {k[7:]: p for k, p in checkpoint["model"].items()}
backbone.load_state_dict(state_dict)
num_channels = 512 if name_map[name][0] in ("resnet18", "resnet34") else 2048
super().__init__(backbone, train_backbone, num_channels, return_interm_layers)
def replace_bn(m, name=""):
for attr_str in dir(m):
target_attr = getattr(m, attr_str)
if isinstance(target_attr, torch.nn.BatchNorm2d):
frozen = FrozenBatchNorm2d(target_attr.num_features)
bn = getattr(m, attr_str)
frozen.weight.data.copy_(bn.weight)
frozen.bias.data.copy_(bn.bias)
frozen.running_mean.data.copy_(bn.running_mean)
frozen.running_var.data.copy_(bn.running_var)
setattr(m, attr_str, frozen)
for n, ch in m.named_children():
replace_bn(ch, n)
class GN_8(nn.Module):
def __init__(self, num_channels):
super().__init__()
self.gn = torch.nn.GroupNorm(8, num_channels)
def forward(self, x):
return self.gn(x)
class TimmBackbone(nn.Module):
def __init__(self, name, return_interm_layers, main_layer=-1, group_norm=False):
super().__init__()
backbone = create_model(name, pretrained=True, in_chans=3, features_only=True, out_indices=(1, 2, 3, 4))
with torch.no_grad():
replace_bn(backbone)
num_channels = backbone.feature_info.channels()[-1]
self.body = backbone
self.num_channels = num_channels
self.interm = return_interm_layers
self.main_layer = main_layer
def forward(self, tensor_list):
xs = self.body(tensor_list.tensors)
if not self.interm:
xs = [xs[self.main_layer]]
out = OrderedDict()
for i, x in enumerate(xs):
mask = F.interpolate(tensor_list.mask[None].float(), size=x.shape[-2:]).bool()[0]
out[f"layer{i}"] = NestedTensor(x, mask)
return out
|