Spaces:
Sleeping
Sleeping
File size: 2,557 Bytes
570db9a |
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 |
from torch import nn, Tensor
import torch.nn.functional as F
from torch.hub import load_state_dict_from_url
from typing import Optional
from ..utils import make_vgg_layers, vgg_cfgs, vgg_urls
class VGG(nn.Module):
def __init__(
self,
features: nn.Module,
reduction: Optional[int] = None,
) -> None:
super().__init__()
self.features = features
self.encoder_reduction = 16
self.reduction = self.encoder_reduction if reduction is None else reduction
self.channels = 512
def forward(self, x: Tensor) -> Tensor:
x = self.features(x)
if self.encoder_reduction != self.reduction:
x = F.interpolate(x, scale_factor=self.encoder_reduction / self.reduction, mode="bilinear")
return x
def _load_weights(model: VGG, url: str) -> VGG:
state_dict = load_state_dict_from_url(url)
missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False)
print("Loading pre-trained weights")
if len(missing_keys) > 0:
print(f"Missing keys: {missing_keys}")
if len(unexpected_keys) > 0:
print(f"Unexpected keys: {unexpected_keys}")
return model
def vgg11(reduction: int = 8) -> VGG:
model = VGG(make_vgg_layers(vgg_cfgs["A"]), reduction=reduction)
return _load_weights(model, vgg_urls["vgg11"])
def vgg11_bn(reduction: int = 8) -> VGG:
model = VGG(make_vgg_layers(vgg_cfgs["A"], batch_norm=True), reduction=reduction)
return _load_weights(model, vgg_urls["vgg11_bn"])
def vgg13(reduction: int = 8) -> VGG:
model = VGG(make_vgg_layers(vgg_cfgs["B"]), reduction=reduction)
return _load_weights(model, vgg_urls["vgg13"])
def vgg13_bn(reduction: int = 8) -> VGG:
model = VGG(make_vgg_layers(vgg_cfgs["B"], batch_norm=True), reduction=reduction)
return _load_weights(model, vgg_urls["vgg13_bn"])
def vgg16(reduction: int = 8) -> VGG:
model = VGG(make_vgg_layers(vgg_cfgs["D"]), reduction=reduction)
return _load_weights(model, vgg_urls["vgg16"])
def vgg16_bn(reduction: int = 8) -> VGG:
model = VGG(make_vgg_layers(vgg_cfgs["D"], batch_norm=True), reduction=reduction)
return _load_weights(model, vgg_urls["vgg16_bn"])
def vgg19(reduction: int = 8) -> VGG:
model = VGG(make_vgg_layers(vgg_cfgs["E"]), reduction=reduction)
return _load_weights(model, vgg_urls["vgg19"])
def vgg19_bn(reduction: int = 8) -> VGG:
model = VGG(make_vgg_layers(vgg_cfgs["E"], batch_norm=True), reduction=reduction)
return _load_weights(model, vgg_urls["vgg19_bn"])
|