|
""" pre-trained GAN models for feature visualization |
|
Examples: |
|
# Load a state dict |
|
G = upconvGAN("fc6") |
|
# Load a state dict from a pretrained model |
|
G = upconvGAN("pool5") |
|
G.G.load_state_dict(torch.hub.load_state_dict_from_url(r"https://drive.google.com/uc?export=download&id=1vB_tOoXL064v9D6AKwl0gTs1a7jo68y7",progress=True)) |
|
""" |
|
import os |
|
from os.path import join |
|
from sys import platform |
|
from collections import OrderedDict |
|
import torch |
|
import torch.nn as nn |
|
import torch.nn.functional as F |
|
|
|
model_urls = { |
|
"caffenet": "https://huggingface.co/binxu/DeePSim_DosovitskiyBrox2016/resolve/main/caffenet.pt", |
|
"norm1": "https://huggingface.co/binxu/DeePSim_DosovitskiyBrox2016/resolve/main/upconvGAN_norm1.pt", |
|
"norm2": "https://huggingface.co/binxu/DeePSim_DosovitskiyBrox2016/resolve/main/upconvGAN_norm2.pt", |
|
"conv3": "https://huggingface.co/binxu/DeePSim_DosovitskiyBrox2016/resolve/main/upconvGAN_conv3.pt", |
|
"conv4": "https://huggingface.co/binxu/DeePSim_DosovitskiyBrox2016/resolve/main/upconvGAN_conv4.pt", |
|
"pool5": "https://huggingface.co/binxu/DeePSim_DosovitskiyBrox2016/resolve/main/upconvGAN_pool5.pt", |
|
"fc6": "https://huggingface.co/binxu/DeePSim_DosovitskiyBrox2016/resolve/main/upconvGAN_fc6.pt", |
|
"fc6_eucl": "https://huggingface.co/binxu/DeePSim_DosovitskiyBrox2016/resolve/main/upconvGAN_fc6_eucl.pt", |
|
"fc7": "https://huggingface.co/binxu/DeePSim_DosovitskiyBrox2016/resolve/main/upconvGAN_fc7.pt", |
|
"fc8": "https://huggingface.co/binxu/DeePSim_DosovitskiyBrox2016/resolve/main/upconvGAN_fc8.pt", |
|
} |
|
|
|
def download_file(url, local_path): |
|
import requests |
|
from tqdm.auto import tqdm |
|
response = requests.get(url, stream=True) |
|
file_size = int(response.headers.get('content-length', 0)) |
|
chunk_size = 1024 |
|
progress_bar = tqdm(total=file_size, unit='iB', unit_scale=True) |
|
|
|
with open(local_path, 'wb') as file: |
|
for chunk in response.iter_content(chunk_size=chunk_size): |
|
progress_bar.update(len(chunk)) |
|
file.write(chunk) |
|
|
|
progress_bar.close() |
|
if file_size != 0 and progress_bar.n != file_size: |
|
print("ERROR, something went wrong.") |
|
|
|
|
|
def load_statedict_from_online(name="fc6"): |
|
import requests |
|
from requests.exceptions import HTTPError |
|
torchhome = torch.hub.get_dir() |
|
ckpthome = join(torchhome, "checkpoints") |
|
os.makedirs(ckpthome, exist_ok=True) |
|
filepath = join(ckpthome, "upconvGAN_%s.pt"%name) |
|
if not os.path.exists(filepath): |
|
print(f"downloading weights of `upconvGAN_{name}.pt` to {filepath}") |
|
download_file(model_urls[name], filepath) |
|
|
|
SD = torch.load(filepath) |
|
return SD |
|
|
|
|
|
class View(nn.Module): |
|
def __init__(self, *shape): |
|
super(View, self).__init__() |
|
self.shape = shape |
|
|
|
def forward(self, x): |
|
return x.view(*self.shape) |
|
|
|
|
|
class ScalarMultiply(nn.Module): |
|
def __init__(self, scalar): |
|
super(ScalarMultiply, self).__init__() |
|
self.scalar = scalar |
|
|
|
def forward(self, x): |
|
return self.scalar * x |
|
|
|
|
|
RGB_mean = torch.tensor([123.0, 117.0, 104.0]) |
|
RGB_mean = torch.reshape(RGB_mean, (1, 3, 1, 1)) |
|
class upconvGAN(nn.Module): |
|
def __init__(self, name="fc6", pretrained=True): |
|
super(upconvGAN, self).__init__() |
|
self.name = name |
|
if name == "fc6" or name == "fc7" or name == "fc6_eucl": |
|
self.G = nn.Sequential(OrderedDict([ |
|
('defc7', nn.Linear(in_features=4096, out_features=4096, bias=True)), |
|
('relu_defc7', nn.LeakyReLU(negative_slope=0.3, inplace=True)), |
|
('defc6', nn.Linear(in_features=4096, out_features=4096, bias=True)), |
|
('relu_defc6', nn.LeakyReLU(negative_slope=0.3, inplace=True)), |
|
('defc5', nn.Linear(in_features=4096, out_features=4096, bias=True)), |
|
('relu_defc5', nn.LeakyReLU(negative_slope=0.3, inplace=True)), |
|
('reshape', View((-1, 256, 4, 4))), |
|
('deconv5', nn.ConvTranspose2d(256, 256, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))), |
|
('relu_deconv5', nn.LeakyReLU(negative_slope=0.3, inplace=True)), |
|
('conv5_1', nn.ConvTranspose2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))), |
|
('relu_conv5_1', nn.LeakyReLU(negative_slope=0.3, inplace=True)), |
|
('deconv4', nn.ConvTranspose2d(512, 256, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))), |
|
('relu_deconv4', nn.LeakyReLU(negative_slope=0.3, inplace=True)), |
|
('conv4_1', nn.ConvTranspose2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))), |
|
('relu_conv4_1', nn.LeakyReLU(negative_slope=0.3, inplace=True)), |
|
('deconv3', nn.ConvTranspose2d(256, 128, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))), |
|
('relu_deconv3', nn.LeakyReLU(negative_slope=0.3, inplace=True)), |
|
('conv3_1', nn.ConvTranspose2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))), |
|
('relu_conv3_1', nn.LeakyReLU(negative_slope=0.3, inplace=True)), |
|
('deconv2', nn.ConvTranspose2d(128, 64, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))), |
|
('relu_deconv2', nn.LeakyReLU(negative_slope=0.3, inplace=True)), |
|
('deconv1', nn.ConvTranspose2d(64, 32, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))), |
|
('relu_deconv1', nn.LeakyReLU(negative_slope=0.3, inplace=True)), |
|
('deconv0', nn.ConvTranspose2d(32, 3, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))), |
|
])) |
|
self.codelen = self.G[0].in_features |
|
self.latent_shape = (4096,) |
|
elif name == "fc8": |
|
self.G = nn.Sequential(OrderedDict([ |
|
("defc7", nn.Linear(in_features=1000, out_features=4096, bias=True)), |
|
("relu_defc7", nn.LeakyReLU(negative_slope=0.3, inplace=True)), |
|
("defc6", nn.Linear(in_features=4096, out_features=4096, bias=True)), |
|
("relu_defc6", nn.LeakyReLU(negative_slope=0.3, inplace=True)), |
|
("defc5", nn.Linear(in_features=4096, out_features=4096, bias=True)), |
|
("relu_defc5", nn.LeakyReLU(negative_slope=0.3, inplace=True)), |
|
("reshape", View((-1, 256, 4, 4))), |
|
("deconv5", nn.ConvTranspose2d(256, 256, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))), |
|
("relu_deconv5", nn.LeakyReLU(negative_slope=0.3, inplace=True)), |
|
("conv5_1", nn.ConvTranspose2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))), |
|
("relu_conv5_1", nn.LeakyReLU(negative_slope=0.3, inplace=True)), |
|
("deconv4", nn.ConvTranspose2d(512, 256, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))), |
|
("relu_deconv4", nn.LeakyReLU(negative_slope=0.3, inplace=True)), |
|
("conv4_1", nn.ConvTranspose2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))), |
|
("relu_conv4_1", nn.LeakyReLU(negative_slope=0.3, inplace=True)), |
|
("deconv3", nn.ConvTranspose2d(256, 128, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))), |
|
("relu_deconv3", nn.LeakyReLU(negative_slope=0.3, inplace=True)), |
|
("conv3_1", nn.ConvTranspose2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))), |
|
("relu_conv3_1", nn.LeakyReLU(negative_slope=0.3, inplace=True)), |
|
("deconv2", nn.ConvTranspose2d(128, 64, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))), |
|
("relu_deconv2", nn.LeakyReLU(negative_slope=0.3, inplace=True)), |
|
("deconv1", nn.ConvTranspose2d(64, 32, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))), |
|
("relu_deconv1", nn.LeakyReLU(negative_slope=0.3, inplace=True)), |
|
("deconv0", nn.ConvTranspose2d(32, 3, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))), |
|
])) |
|
self.codelen = self.G[0].in_features |
|
self.latent_shape = (1000,) |
|
elif name == "pool5": |
|
self.G = nn.Sequential(OrderedDict([ |
|
('Rconv6', nn.Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))), |
|
('Rrelu6', nn.LeakyReLU(negative_slope=0.3, inplace=True)), |
|
('Rconv7', nn.Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))), |
|
('Rrelu7', nn.LeakyReLU(negative_slope=0.3, inplace=True)), |
|
('Rconv8', nn.Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1))), |
|
('Rrelu8', nn.LeakyReLU(negative_slope=0.3, inplace=True)), |
|
('deconv5', nn.ConvTranspose2d(512, 256, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))), |
|
('relu_deconv5', nn.LeakyReLU(negative_slope=0.3, inplace=True)), |
|
('conv5_1', nn.ConvTranspose2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))), |
|
('relu_conv5_1', nn.LeakyReLU(negative_slope=0.3, inplace=True)), |
|
('deconv4', nn.ConvTranspose2d(512, 256, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))), |
|
('relu_deconv4', nn.LeakyReLU(negative_slope=0.3, inplace=True)), |
|
('conv4_1', nn.ConvTranspose2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))), |
|
('relu_conv4_1', nn.LeakyReLU(negative_slope=0.3, inplace=True)), |
|
('deconv3', nn.ConvTranspose2d(256, 128, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))), |
|
('relu_deconv3', nn.LeakyReLU(negative_slope=0.3, inplace=True)), |
|
('conv3_1', nn.ConvTranspose2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))), |
|
('relu_conv3_1', nn.LeakyReLU(negative_slope=0.3, inplace=True)), |
|
('deconv2', nn.ConvTranspose2d(128, 64, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))), |
|
('relu_deconv2', nn.LeakyReLU(negative_slope=0.3, inplace=True)), |
|
('deconv1', nn.ConvTranspose2d(64, 32, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))), |
|
('relu_deconv1', nn.LeakyReLU(negative_slope=0.3, inplace=True)), |
|
('deconv0', nn.ConvTranspose2d(32, 3, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))), ])) |
|
self.codelen = self.G[0].in_channels |
|
self.latent_shape = (256, 6, 6) |
|
elif name == "conv4": |
|
self.G = nn.Sequential(OrderedDict([ |
|
('Rconv6', nn.Conv2d(384, 384, kernel_size=(3, 3), stride=(1, 1))), |
|
('Rrelu6', nn.LeakyReLU(negative_slope=0.3, inplace=True)), |
|
('Rconv7', nn.Conv2d(384, 512, kernel_size=(3, 3), stride=(1, 1))), |
|
('Rrelu7', nn.LeakyReLU(negative_slope=0.3, inplace=True)), |
|
('Rconv8', nn.Conv2d(512, 512, kernel_size=(2, 2), stride=(1, 1))), |
|
('Rrelu8', nn.LeakyReLU(negative_slope=0.3, inplace=True)), |
|
('deconv5', nn.ConvTranspose2d(512, 256, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))), |
|
('relu_deconv5', nn.LeakyReLU(negative_slope=0.3, inplace=True)), |
|
('conv5_1', nn.ConvTranspose2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))), |
|
('relu_conv5_1', nn.LeakyReLU(negative_slope=0.3, inplace=True)), |
|
('deconv4', nn.ConvTranspose2d(256, 128, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))), |
|
('relu_deconv4', nn.LeakyReLU(negative_slope=0.3, inplace=True)), |
|
('conv4_1', nn.ConvTranspose2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))), |
|
('relu_conv4_1', nn.LeakyReLU(negative_slope=0.3, inplace=True)), |
|
('deconv3', nn.ConvTranspose2d(128, 128, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))), |
|
('relu_deconv3', nn.LeakyReLU(negative_slope=0.3, inplace=True)), |
|
('conv3_1', nn.ConvTranspose2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))), |
|
('relu_conv3_1', nn.LeakyReLU(negative_slope=0.3, inplace=True)), |
|
('deconv2', nn.ConvTranspose2d(128, 64, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))), |
|
('relu_deconv2', nn.LeakyReLU(negative_slope=0.3, inplace=True)), |
|
('conv2_1', nn.Conv2d(64, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))), |
|
('relu_conv2_1', nn.LeakyReLU(negative_slope=0.3, inplace=True)), |
|
('deconv1', nn.ConvTranspose2d(32, 16, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))), |
|
('relu_deconv1', nn.LeakyReLU(negative_slope=0.3, inplace=True)), |
|
('conv1_1', nn.Conv2d(16, 3, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))), |
|
('tanh', nn.Tanh()), |
|
("scaling", ScalarMultiply(255.0)), |
|
])) |
|
self.codelen = self.G[0].in_channels |
|
self.latent_shape = (384, 13, 13) |
|
elif name == "conv3": |
|
self.G = nn.Sequential(OrderedDict([ |
|
('Rconv6', nn.Conv2d(384, 384, kernel_size=(3, 3), stride=(1, 1))), |
|
('Rrelu6', nn.LeakyReLU(negative_slope=0.3, inplace=True)), |
|
('Rconv7', nn.Conv2d(384, 512, kernel_size=(3, 3), stride=(1, 1))), |
|
('Rrelu7', nn.LeakyReLU(negative_slope=0.3, inplace=True)), |
|
('Rconv8', nn.Conv2d(512, 512, kernel_size=(2, 2), stride=(1, 1))), |
|
('Rrelu8', nn.LeakyReLU(negative_slope=0.3, inplace=True)), |
|
('deconv5', nn.ConvTranspose2d(512, 256, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))), |
|
('relu_deconv5', nn.LeakyReLU(negative_slope=0.3, inplace=True)), |
|
('conv5_1', nn.ConvTranspose2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))), |
|
('relu_conv5_1', nn.LeakyReLU(negative_slope=0.3, inplace=True)), |
|
('deconv4', nn.ConvTranspose2d(256, 128, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))), |
|
('relu_deconv4', nn.LeakyReLU(negative_slope=0.3, inplace=True)), |
|
('conv4_1', nn.ConvTranspose2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))), |
|
('relu_conv4_1', nn.LeakyReLU(negative_slope=0.3, inplace=True)), |
|
('deconv3', nn.ConvTranspose2d(128, 128, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))), |
|
('relu_deconv3', nn.LeakyReLU(negative_slope=0.3, inplace=True)), |
|
('conv3_1', nn.ConvTranspose2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))), |
|
('relu_conv3_1', nn.LeakyReLU(negative_slope=0.3, inplace=True)), |
|
('deconv2', nn.ConvTranspose2d(128, 64, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))), |
|
('relu_deconv2', nn.LeakyReLU(negative_slope=0.3, inplace=True)), |
|
('conv2_1', nn.Conv2d(64, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))), |
|
('relu_conv2_1', nn.LeakyReLU(negative_slope=0.3, inplace=True)), |
|
('deconv1', nn.ConvTranspose2d(32, 16, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))), |
|
('relu_deconv1', nn.LeakyReLU(negative_slope=0.3, inplace=True)), |
|
('conv1_1', nn.Conv2d(16, 3, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))), |
|
('tanh', nn.Tanh()), |
|
("scaling", ScalarMultiply(255.0)), |
|
])) |
|
self.codelen = self.G[0].in_channels |
|
self.latent_shape = (384, 13, 13) |
|
elif name == "norm2": |
|
self.G = nn.Sequential(OrderedDict([ |
|
('Rconv6', nn.Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(2, 2))), |
|
('Rrelu6', nn.LeakyReLU(negative_slope=0.3, inplace=True)), |
|
('Rconv7', nn.Conv2d(256, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))), |
|
('Rrelu7', nn.LeakyReLU(negative_slope=0.3, inplace=True)), |
|
('Rconv8', nn.Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))), |
|
('Rrelu8', nn.LeakyReLU(negative_slope=0.3, inplace=True)), |
|
('deconv4', nn.ConvTranspose2d(128, 128, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))), |
|
('relu_deconv4', nn.LeakyReLU(negative_slope=0.3, inplace=True)), |
|
('conv4_1', nn.Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))), |
|
('relu_conv4_1', nn.LeakyReLU(negative_slope=0.3, inplace=True)), |
|
('deconv3', nn.ConvTranspose2d(128, 64, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))), |
|
('relu_deconv3', nn.LeakyReLU(negative_slope=0.3, inplace=True)), |
|
('conv3_1', nn.Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))), |
|
('relu_conv3_1', nn.LeakyReLU(negative_slope=0.3, inplace=True)), |
|
('deconv2', nn.ConvTranspose2d(64, 32, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))), |
|
('relu_deconv2', nn.LeakyReLU(negative_slope=0.3, inplace=True)), |
|
('conv2_1', nn.Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))), |
|
('relu_conv2_1', nn.LeakyReLU(negative_slope=0.3, inplace=True)), |
|
('deconv1', nn.ConvTranspose2d(32, 16, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))), |
|
('conv1_1', nn.Conv2d(16, 3, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))), |
|
('tanh', nn.Tanh()), |
|
("scaling", ScalarMultiply(255.0)), |
|
])) |
|
self.codelen = self.G[0].in_channels |
|
self.latent_shape = (256, 13, 13) |
|
elif name == "norm1": |
|
self.G = nn.Sequential(OrderedDict([ |
|
('Rconv6', nn.Conv2d(96, 128, kernel_size=(3, 3), stride=(2, 2), padding=(2, 2))), |
|
('Rrelu6', nn.LeakyReLU(negative_slope=0.3, inplace=True)), |
|
('Rconv7', nn.Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))), |
|
('Rrelu7', nn.LeakyReLU(negative_slope=0.3, inplace=True)), |
|
('Rconv8', nn.Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))), |
|
('Rrelu8', nn.LeakyReLU(negative_slope=0.3, inplace=True)), |
|
('deconv4', nn.ConvTranspose2d(128, 128, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))), |
|
('relu_deconv4', nn.LeakyReLU(negative_slope=0.3, inplace=True)), |
|
('conv4_1', nn.Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))), |
|
('relu_conv4_1', nn.LeakyReLU(negative_slope=0.3, inplace=True)), |
|
('deconv3', nn.ConvTranspose2d(128, 64, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))), |
|
('relu_deconv3', nn.LeakyReLU(negative_slope=0.3, inplace=True)), |
|
('conv3_1', nn.Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))), |
|
('relu_conv3_1', nn.LeakyReLU(negative_slope=0.3, inplace=True)), |
|
('deconv2', nn.ConvTranspose2d(64, 32, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))), |
|
('relu_deconv2', nn.LeakyReLU(negative_slope=0.3, inplace=True)), |
|
('conv2_1', nn.Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))), |
|
('relu_conv2_1', nn.LeakyReLU(negative_slope=0.3, inplace=True)), |
|
('deconv1', nn.ConvTranspose2d(32, 16, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))), |
|
('conv1_1', nn.Conv2d(16, 3, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))), |
|
('tanh', nn.Tanh()), |
|
("scaling", ScalarMultiply(255.0)), |
|
])) |
|
self.codelen = self.G[0].in_channels |
|
self.latent_shape = (96, 27, 27) |
|
|
|
if pretrained: |
|
SD = load_statedict_from_online(name) |
|
SDnew = OrderedDict() |
|
for name, W in SD.items(): |
|
name = name.replace(".1.", ".") |
|
SDnew[name] = W |
|
self.G.load_state_dict(SDnew) |
|
|
|
def sample_vector(self, sampn=1, device="cuda", noise_std=1.0): |
|
return torch.randn(sampn, self.codelen, device=device) * noise_std |
|
|
|
def forward(self, x): |
|
return self.G(x)[:, [2, 1, 0], :, :] |
|
|
|
def visualize(self, x, scale=1.0): |
|
raw = self.G(x)[:, [2, 1, 0], :, :] |
|
return torch.clamp(raw + RGB_mean.to(raw.device), 0, 255.0) / 255.0 * scale |
|
|
|
def visualize_batch(self, x_arr, scale=1.0, B=42, ): |
|
coden = x_arr.shape[0] |
|
img_all = [] |
|
csr = 0 |
|
with torch.no_grad(): |
|
while csr < coden: |
|
csr_end = min(csr + B, coden) |
|
imgs = self.visualize(x_arr[csr:csr_end, :].cuda(), scale).cpu() |
|
img_all.append(imgs) |
|
csr = csr_end |
|
img_all = torch.cat(img_all, dim=0) |
|
return img_all |
|
|
|
def render(self, x, scale=1.0, B=42): |
|
coden = x.shape[0] |
|
img_all = [] |
|
csr = 0 |
|
while csr < coden: |
|
csr_end = min(csr + B, coden) |
|
with torch.no_grad(): |
|
imgs = self.visualize(torch.from_numpy(x[csr:csr_end, :]).float().cuda(), scale).permute(2,3,1,0).cpu().numpy() |
|
img_all.extend([imgs[:, :, :, imgi] for imgi in range(imgs.shape[3])]) |
|
csr = csr_end |
|
return img_all |
|
|
|
def visualize_batch_np(self, codes_all_arr, scale=1.0, B=42, verbose=False): |
|
coden = codes_all_arr.shape[0] |
|
img_all = None |
|
csr = 0 |
|
with torch.no_grad(): |
|
while csr < coden: |
|
csr_end = min(csr + B, coden) |
|
imgs = self.visualize(torch.from_numpy(codes_all_arr[csr:csr_end, :]).float().cuda(), scale).cpu() |
|
img_all = imgs if img_all is None else torch.cat((img_all, imgs), dim=0) |
|
csr = csr_end |
|
return img_all |
|
|
|
|
|
class Caffenet(nn.Module): |
|
def __init__(self, pretrained=True, ): |
|
super(Caffenet, self).__init__() |
|
self.net = nn.Sequential(OrderedDict([ |
|
|
|
('conv1', nn.Conv2d(3, 96, kernel_size=(11, 11), stride=(4, 4))), |
|
('relu1', nn.ReLU(inplace=True)), |
|
('pool1', nn.MaxPool2d(kernel_size=3, stride=2, padding=0, dilation=1, ceil_mode=True)), |
|
('norm1', nn.LocalResponseNorm(size=5, alpha=0.0001, beta=0.75, k=1.0)), |
|
|
|
('conv2', nn.Conv2d(96, 256, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2), groups=2)), |
|
('relu2', nn.ReLU(inplace=True)), |
|
('pool2', nn.MaxPool2d(kernel_size=3, stride=2, padding=0, dilation=1, ceil_mode=True)), |
|
('norm2', nn.LocalResponseNorm(size=5, alpha=0.0001, beta=0.75, k=1.0)), |
|
|
|
('conv3', nn.Conv2d(256, 384, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))), |
|
('relu3', nn.ReLU(inplace=True)), |
|
|
|
('conv4', nn.Conv2d(384, 384, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=2)), |
|
('relu4', nn.ReLU(inplace=True)), |
|
|
|
('conv5', nn.Conv2d(384, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=2)), |
|
('relu5', nn.ReLU(inplace=True)), |
|
('pool5', nn.MaxPool2d(kernel_size=3, stride=2, padding=0, dilation=1, ceil_mode=True)), |
|
|
|
('flatten', nn.Flatten()), |
|
('fc6', nn.Linear(9216, 4096)), |
|
('relu6', nn.ReLU(inplace=True)), |
|
('fc7', nn.Linear(4096, 4096)), |
|
('relu7', nn.ReLU(inplace=True)), |
|
('fc8', nn.Linear(4096, 1000)) |
|
])) |
|
if pretrained: |
|
SD = load_statedict_from_online("caffenet") |
|
SDnew = OrderedDict() |
|
for name, W in SD.items(): |
|
name = name.replace(".1.", ".") |
|
SDnew[name] = W |
|
self.net.load_state_dict(SDnew) |
|
|
|
def forward(self, x, preproc=False, scale=1.0): |
|
if preproc: |
|
x = x.float() / scale * 255.0 |
|
x = x - RGB_mean.to(x.device) |
|
x = x[:, [2, 1, 0], :, :] |
|
return self.net(x) |
|
|
|
|
