ai-art

 1import torch
 2import torch.nn as nn
 3import torch.nn.functional as F
 4
 5from torchvision.models import VGG
 6try:
 7    from torch.hub import load_state_dict_from_url
 8except ImportError:
 9    from torch.utils.model_zoo import load_url as load_state_dict_from_url
10
11from typing import List, Union, cast
12
13
14def adopt_weight(weight, global_step, threshold=0, value=0.):
15    if global_step < threshold:
16        weight = value
17    return weight
18
19
20def hinge_d_loss(logits_real, logits_fake):
21    loss_real = torch.mean(F.relu(1. - logits_real))
22    loss_fake = torch.mean(F.relu(1. + logits_fake))
23    d_loss = 0.5 * (loss_real + loss_fake)
24    return d_loss
25
26
27def vanilla_d_loss(logits_real, logits_fake):
28    d_loss = 0.5 * (
29        torch.mean(torch.nn.functional.softplus(-logits_real)) +
30        torch.mean(torch.nn.functional.softplus(logits_fake)))
31    return d_loss
32
33
34def normalize_tensor(x, eps=1e-10):
35    norm_factor = torch.sqrt(torch.sum(x**2, dim=1, keepdim=True))
36    return x / (norm_factor + eps)
37
38
39def spatial_average(x, keepdim=True):
40    return x.mean([2, 3], keepdim=keepdim)
41
42
43def make_layers(cfg: List[Union[str, int]], batch_norm: bool = False) -> nn.Sequential:
44    layers: List[nn.Module] = []
45    in_channels = 3
46    for v in cfg:
47        if v == 'M':
48            layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
49        else:
50            v = cast(int, v)
51            conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=1)
52            if batch_norm:
53                layers += [conv2d, nn.BatchNorm2d(v), nn.ReLU(inplace=True)]
54            else:
55                layers += [conv2d, nn.ReLU(inplace=True)]
56            in_channels = v
57    return nn.Sequential(*layers)
58
59
60def load_vgg(model_dir: str, pretrained: bool = False, **kwargs):
61    if pretrained:
62        kwargs['init_weights'] = False
63
64    cfg = [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M', 512, 512, 512, 'M']
65    model = VGG(make_layers(cfg, batch_norm=False), **kwargs)
66
67    if pretrained:
68        state_dict = load_state_dict_from_url('https://download.pytorch.org/models/vgg16-397923af.pth',
69                                              model_dir=model_dir,
70                                              file_name="vgg16-397923af.pth",
71                                              progress=True)
72        model.load_state_dict(state_dict)
73        print(f"Loaded pretrained VGG16 model from '{model_dir}/vgg16-397923af.pth'")
74
75    return model