escapist413/StyleFusion

1import torch.nn as nn2from torchvision import models3from torchvision.models import VGG19_Weights4 5 6class VGG(nn.Module):7    def __init__(self, content_layers, style_layers):8        super(VGG, self).__init__()9        self.model = models.vgg19(weights=VGG19_Weights.IMAGENET1K_V1).features10        self.content_layers = content_layers.keys()11        self.style_layers = style_layers.keys()12 13        # 冻结模型的所有参数14        for param in self.model.parameters():15            param.requires_grad = False16 17    def forward(self, x):18        """19        对vgg19网络的包装，前向传播时保留了内容层和风格层的中间输出20        :param x:21        :return: 内容层和风格层的特征图22        """23        content_features = {}24        style_features = {}25 26        for name, layer in self.model._modules.items():27            x = layer(x)28            if name in self.content_layers:29                content_features[name] = x30            if name in self.style_layers:31                style_features[name] = x32 33        return content_features, style_features34 35 36class ResBlock(nn.Module):37 38    def __init__(self, c):39        super(ResBlock, self).__init__()40        self.layer = nn.Sequential(41            nn.Conv2d(c, c, 3, 1, 1, bias=False),42            nn.InstanceNorm2d(c),43            nn.ReLU(True),44            nn.Conv2d(c, c, 3, 1, 1, bias=False),45            nn.InstanceNorm2d(c)46        )47 48    def forward(self, x):49        return x + self.layer(x)50 51 52class TransNet(nn.Module):53    def __init__(self, input_size):54        """55        实时内容生成网络56        """57        super(TransNet, self).__init__()58        self.input_size = input_size59        self.layer = nn.Sequential(60            ###################下采样层################61            nn.Conv2d(in_channels=3, out_channels=32, kernel_size=9, stride=1, padding=4, bias=False),62            nn.InstanceNorm2d(32),63            nn.ReLU(True),64            nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, stride=2, padding=1, bias=False),65            nn.InstanceNorm2d(64),66            nn.ReLU(True),67            nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, stride=2, padding=1, bias=False),68            nn.InstanceNorm2d(128),69            nn.ReLU(True),70 71            ##################残差层##################72            ResBlock(128),73            ResBlock(128),74            ResBlock(128),75            ResBlock(128),76            ResBlock(128),77 78            ################上采样层##################79            nn.Upsample(scale_factor=2, mode='nearest'),80            nn.Conv2d(in_channels=128, out_channels=64, kernel_size=3, stride=1, padding=1, bias=False),81            nn.InstanceNorm2d(64),82            nn.ReLU(True),83            nn.Upsample(scale_factor=2, mode='nearest'),84            nn.Conv2d(in_channels=64, out_channels=32, kernel_size=3, stride=1, padding=1, bias=False),85            nn.InstanceNorm2d(32),86            nn.ReLU(True),87 88            ###############输出层#####################89            nn.Conv2d(in_channels=32, out_channels=3, kernel_size=9, stride=1, padding=4, bias=False),90            nn.Sigmoid(),91            nn.AdaptiveAvgPool2d(self.input_size)92        )93 94    def forward(self, x):95        return self.layer(x)96