You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
88 lines
2.9 KiB
88 lines
2.9 KiB
import paddle |
|
import paddle.nn as nn |
|
import numpy as np |
|
from ppgan.utils.download import get_path_from_url |
|
|
|
model_urls = { |
|
'caffevgg19': ('https://paddlegan.bj.bcebos.com/models/vgg19_no_fc.npy', |
|
'8ea1ef2374f8684b6cea9f300849be81') |
|
} |
|
|
|
|
|
class CaffeVGG19(nn.Layer): |
|
cfg = [ |
|
64, 64, 'M', 128, 128, 'M', 256, 256, 256, 256, 'M', 512, 512, 512, 512, |
|
'M', 512, 512, 512, 512, 'M' |
|
] |
|
|
|
def __init__(self, output_index: int=26) -> None: |
|
super().__init__() |
|
arch = 'caffevgg19' |
|
weights_path = get_path_from_url(model_urls[arch][0], |
|
model_urls[arch][1]) |
|
data_dict: dict = np.load( |
|
weights_path, encoding='latin1', allow_pickle=True).item() |
|
self.features = self.make_layers(self.cfg, data_dict) |
|
del data_dict |
|
self.features = nn.Sequential(*self.features.sublayers()[:output_index]) |
|
mean = paddle.to_tensor([103.939, 116.779, 123.68]) |
|
self.mean = mean.unsqueeze(0).unsqueeze(-1).unsqueeze(-1) |
|
|
|
def _process(self, x): |
|
# value to 255 |
|
rgb = (x * 0.5 + 0.5) * 255 |
|
# rgb to bgr |
|
bgr = paddle.stack((rgb[:, 2, :, :], rgb[:, 1, :, :], rgb[:, 0, :, :]), |
|
1) |
|
# vgg norm |
|
return bgr - self.mean |
|
|
|
def _forward_impl(self, x): |
|
x = self._process(x) |
|
# NOTE get output with out relu activation |
|
x = self.features(x) |
|
return x |
|
|
|
def forward(self, x): |
|
return self._forward_impl(x) |
|
|
|
@staticmethod |
|
def get_conv_filter(data_dict, name): |
|
return data_dict[name][0] |
|
|
|
@staticmethod |
|
def get_bias(data_dict, name): |
|
return data_dict[name][1] |
|
|
|
@staticmethod |
|
def get_fc_weight(data_dict, name): |
|
return data_dict[name][0] |
|
|
|
def make_layers(self, cfg, data_dict, batch_norm=False) -> nn.Sequential: |
|
layers = [] |
|
in_channels = 3 |
|
block = 1 |
|
number = 1 |
|
for v in cfg: |
|
if v == 'M': |
|
layers += [nn.MaxPool2D(kernel_size=2, stride=2)] |
|
block += 1 |
|
number = 1 |
|
else: |
|
conv2d = nn.Conv2D(in_channels, v, kernel_size=3, padding=1) |
|
""" set value """ |
|
weight = paddle.to_tensor( |
|
self.get_conv_filter(data_dict, f'conv{block}_{number}')) |
|
weight = weight.transpose((3, 2, 0, 1)) |
|
bias = paddle.to_tensor( |
|
self.get_bias(data_dict, f'conv{block}_{number}')) |
|
conv2d.weight.set_value(weight) |
|
conv2d.bias.set_value(bias) |
|
number += 1 |
|
if batch_norm: |
|
layers += [conv2d, nn.BatchNorm2D(v), nn.ReLU()] |
|
else: |
|
layers += [conv2d, nn.ReLU()] |
|
in_channels = v |
|
|
|
return nn.Sequential(*layers)
|
|
|