SPADE 代码略解 ade20k数据集

2024-03-06 20:38

本文主要是介绍SPADE 代码略解 ade20k数据集,希望对大家解决编程问题提供一定的参考价值,需要的开发者们随着小编来一起学习吧!

对paper《Semantic Image Synthesis with Spatially-Adaptive Normalization》的代码梳理,主要用的是ade20k数据集

train.py

        # Training# train generatorif i % opt.D_steps_per_G == 0:trainer.run_generator_one_step(data_i)

通过这一行代码开始训练。

trainers/pix2pix_trainer.py

    def run_generator_one_step(self, data):self.optimizer_G.zero_grad()g_losses, generated = self.pix2pix_model(data, mode='generator')g_loss = sum(g_losses.values()).mean()g_loss.backward()self.optimizer_G.step()self.g_losses = g_lossesself.generated = generated

在这个函数中通过调用self.pix2pix_model(data,mode='generator')训练。(我在想为什么这样的模型架构都要用pix2pix命名,后来发现原来这种从一张图像生成另一张图像的架构都叫做pix2pix,和cgan的区别在于cgan是从噪音+图像生成另一张图像的)

class Pix2PixTrainer():"""Trainer creates the model and optimizers, and uses them toupdates the weights of the network while reporting lossesand the latest visuals to visualize the progress in training."""def __init__(self, opt):self.opt = optself.pix2pix_model = Pix2PixModel(opt)if len(opt.gpu_ids) > 0:self.pix2pix_model = DataParallelWithCallback(self.pix2pix_model,device_ids=opt.gpu_ids)self.pix2pix_model_on_one_gpu = self.pix2pix_model.moduleelse:self.pix2pix_model_on_one_gpu = self.pix2pix_modelself.generated = Noneif opt.isTrain:self.optimizer_G, self.optimizer_D = \self.pix2pix_model_on_one_gpu.create_optimizers(opt)self.old_lr = opt.lr

在这个类的构造函数里定义了self.pix2pix_model是通过Pix2PixModel来的,其中多卡的时候有个跨卡BN的操作。

models/pix2pix_model.py

    def forward(self, data, mode):input_semantics, real_image = self.preprocess_input(data)if mode == 'generator':g_loss, generated = self.compute_generator_loss(input_semantics, real_image)return g_loss, generatedelif mode == 'discriminator':d_loss = self.compute_discriminator_loss(input_semantics, real_image)return d_losselif mode == 'encode_only':z, mu, logvar = self.encode_z(real_image)return mu, logvarelif mode == 'inference':with torch.no_grad():fake_image, _ = self.generate_fake(input_semantics, real_image)return fake_imageelse:raise ValueError("|mode| is invalid")

当mode为generator的时候,调用self.compute_generator_loss(input_semantics,real_image)

这里看下这两个输入条件是通过self.preprocess_input(data)来的,具体是:

    def preprocess_input(self, data):# move to GPU and change data typesdata['label'] = data['label'].long()if self.use_gpu():data['label'] = data['label'].cuda()data['instance'] = data['instance'].cuda()data['image'] = data['image'].cuda()# create one-hot label maplabel_map = data['label']bs, _, h, w = label_map.size()nc = self.opt.label_nc + 1 if self.opt.contain_dontcare_label \else self.opt.label_ncinput_label = self.FloatTensor(bs, nc, h, w).zero_()input_semantics = input_label.scatter_(1, label_map, 1.0)# concatenate instance map if it existsif not self.opt.no_instance:inst_map = data['instance']instance_edge_map = self.get_edges(inst_map)input_semantics = torch.cat((input_semantics, instance_edge_map), dim=1)return input_semantics, data['image']

    def compute_generator_loss(self, input_semantics, real_image):G_losses = {}fake_image, KLD_loss = self.generate_fake(input_semantics, real_image, compute_kld_loss=self.opt.use_vae)###在这里输入数据,生成fake_imageif self.opt.use_vae:G_losses['KLD'] = KLD_losspred_fake, pred_real = self.discriminate(input_semantics, fake_image, real_image)G_losses['GAN'] = self.criterionGAN(pred_fake, True,for_discriminator=False)if not self.opt.no_ganFeat_loss:num_D = len(pred_fake)GAN_Feat_loss = self.FloatTensor(1).fill_(0)for i in range(num_D):  # for each discriminator# last output is the final prediction, so we exclude itnum_intermediate_outputs = len(pred_fake[i]) - 1for j in range(num_intermediate_outputs):  # for each layer outputunweighted_loss = self.criterionFeat(pred_fake[i][j], pred_real[i][j].detach())GAN_Feat_loss += unweighted_loss * self.opt.lambda_feat / num_DG_losses['GAN_Feat'] = GAN_Feat_lossif not self.opt.no_vgg_loss:G_losses['VGG'] = self.criterionVGG(fake_image, real_image) \* self.opt.lambda_vggreturn G_losses, fake_image
    def generate_fake(self, input_semantics, real_image, compute_kld_loss=False):z = NoneKLD_loss = Noneif self.opt.use_vae:z, mu, logvar = self.encode_z(real_image)if compute_kld_loss:KLD_loss = self.KLDLoss(mu, logvar) * self.opt.lambda_kldfake_image = self.netG(input_semantics, z=z)assert (not compute_kld_loss) or self.opt.use_vae, \"You cannot compute KLD loss if opt.use_vae == False"return fake_image, KLD_loss
#在训练时,采用的是不考虑vae的,也就是不会对real_image做encode操作得到一个z,
#此处的z为None,送去生成网络self.netG的只有语义标签图和Z,没有real_image。
#我一开始很迷惑这一步,只有纯语义标签用来生成,这样的话mask怎么去学习image的风格呢
#先往下看
class Pix2PixModel(torch.nn.Module):@staticmethoddef modify_commandline_options(parser, is_train):networks.modify_commandline_options(parser, is_train)return parserdef __init__(self, opt):super().__init__()self.opt = optself.FloatTensor = torch.cuda.FloatTensor if self.use_gpu() \else torch.FloatTensorself.ByteTensor = torch.cuda.ByteTensor if self.use_gpu() \else torch.ByteTensorself.netG, self.netD, self.netE = self.initialize_networks(opt) 
##在这里得到初始化网络后(这里的初始化不是真的在做初始化)的self.netG,
    def initialize_networks(self, opt):netG = networks.define_G(opt) ##在这里得到netGnetD = networks.define_D(opt) if opt.isTrain else NonenetE = networks.define_E(opt) if opt.use_vae else Noneif not opt.isTrain or opt.continue_train:netG = util.load_network(netG, 'G', opt.which_epoch, opt)if opt.isTrain:netD = util.load_network(netD, 'D', opt.which_epoch, opt)if opt.use_vae:netE = util.load_network(netE, 'E', opt.which_epoch, opt)return netG, netD, netE

models/networks/__init__.py (是network下面的__init__.py而不是model下面的)

def define_G(opt):netG_cls = find_network_using_name(opt.netG, 'generator')return create_network(netG_cls, opt)
def find_network_using_name(target_network_name, filename):
### target_network_name 是SPADE,filename是generatortarget_class_name = target_network_name + filenamemodule_name = 'models.networks.' + filenamenetwork = util.find_class_in_module(target_class_name, module_name)#在models.networks.generator里面找到SPADEGenerator这个模块并返回assert issubclass(network, BaseNetwork), \"Class %s should be a subclass of BaseNetwork" % networkreturn network
def create_network(cls, opt):net = cls(opt) #输入一些网络参数net.print_network() #打印网络if len(opt.gpu_ids) > 0:assert(torch.cuda.is_available())net.cuda()net.init_weights(opt.init_type, opt.init_variance) #这里才是真的在做初始化网络return net

接下来,具体看它调用的SPADEGenerator的网络结构

models/networks/generator.py

class SPADEGenerator(BaseNetwork):@staticmethoddef modify_commandline_options(parser, is_train):parser.set_defaults(norm_G='spectralspadesyncbatch3x3')parser.add_argument('--num_upsampling_layers',choices=('normal', 'more', 'most'), default='normal',help="If 'more', adds upsampling layer between the two middle resnet blocks. If 'most', also add one more upsampling + resnet layer at the end of the generator")return parserdef __init__(self, opt):super().__init__()self.opt = optnf = opt.ngfself.sw, self.sh = self.compute_latent_vector_size(opt)### 输入为256x256,得到的sw=2,sh=2 计算潜向量的大小if opt.use_vae:# In case of VAE, we will sample from random z vectorself.fc = nn.Linear(opt.z_dim, 16 * nf * self.sw * self.sh)else:# Otherwise, we make the network deterministic by starting with# downsampled segmentation map instead of random zself.fc = nn.Conv2d(self.opt.semantic_nc, 16 * nf, 3, padding=1)self.head_0 = SPADEResnetBlock(16 * nf, 16 * nf, opt)self.G_middle_0 = SPADEResnetBlock(16 * nf, 16 * nf, opt)self.G_middle_1 = SPADEResnetBlock(16 * nf, 16 * nf, opt)self.up_0 = SPADEResnetBlock(16 * nf, 8 * nf, opt)self.up_1 = SPADEResnetBlock(8 * nf, 4 * nf, opt)self.up_2 = SPADEResnetBlock(4 * nf, 2 * nf, opt)self.up_3 = SPADEResnetBlock(2 * nf, 1 * nf, opt)final_nc = nfif opt.num_upsampling_layers == 'most':self.up_4 = SPADEResnetBlock(1 * nf, nf // 2, opt)final_nc = nf // 2self.conv_img = nn.Conv2d(final_nc, 3, 3, padding=1)self.up = nn.Upsample(scale_factor=2)def compute_latent_vector_size(self, opt):if opt.num_upsampling_layers == 'normal':num_up_layers = 5elif opt.num_upsampling_layers == 'more':num_up_layers = 6elif opt.num_upsampling_layers == 'most':num_up_layers = 7else:raise ValueError('opt.num_upsampling_layers [%s] not recognized' %opt.num_upsampling_layers)sw = opt.crop_size // (2**num_up_layers)sh = round(sw / opt.aspect_ratio)return sw, shdef forward(self, input, z=None):seg = input ### 这里的input是语义标签图if self.opt.use_vae:# we sample z from unit normal and reshape the tensorif z is None:z = torch.randn(input.size(0), self.opt.z_dim,dtype=torch.float32, device=input.get_device())x = self.fc(z)x = x.view(-1, 16 * self.opt.ngf, self.sh, self.sw)else:# we downsample segmap and run convolutionx = F.interpolate(seg, size=(self.sh, self.sw)) ##对语义标签图插值后变成size更小的特征图?x = self.fc(x) #卷积操作x = self.head_0(x, seg) #这里的x已经变成了sh X sw这么大,通道为16*nf的特征图了,而seg还是原图大小,特征通道为151的初始inputx = self.up(x) #上采样2倍x = self.G_middle_0(x, seg) #不改变通道值的SPADEResnetBlock,建议先去看一下SPADEResnetBlock的构造if self.opt.num_upsampling_layers == 'more' or \self.opt.num_upsampling_layers == 'most':x = self.up(x) x = self.G_middle_1(x, seg) #SPADEResnetBlockx = self.up(x)x = self.up_0(x, seg)x = self.up(x)x = self.up_1(x, seg)x = self.up(x)x = self.up_2(x, seg)x = self.up(x)x = self.up_3(x, seg)if self.opt.num_upsampling_layers == 'most':x = self.up(x)x = self.up_4(x, seg)x = self.conv_img(F.leaky_relu(x, 2e-1))x = F.tanh(x)return x

这一步我觉得需要注意的是输入到generator的input,把mask作为input是为了得到spatial信息的。但我之前一直以为是把mask做encode之后用image来学习仿射变换的参数“注射”到特征图的标准化中,原来generator从头到尾都用不到image啊,估计只有loss的时候才用到。这里提出的生成器里,主要分为1.用vae(这里又分为有没有提供real image)2.不用vae 。用vae的时候如果提供了real image,就算real image的均值和方差得到一个z向量,如果没有提供,就生成一个符合标准正太分布的随机噪声,然后连接全连接层生成一个z向量。不用vae的时候是对segmantic map做降采样处理作为输入。

models/networks/architecture.py

class SPADEResnetBlock(nn.Module):def __init__(self, fin, fout, opt):super().__init__()# Attributesself.learned_shortcut = (fin != fout)fmiddle = min(fin, fout)# create conv layersself.conv_0 = nn.Conv2d(fin, fmiddle, kernel_size=3, padding=1)self.conv_1 = nn.Conv2d(fmiddle, fout, kernel_size=3, padding=1)if self.learned_shortcut:self.conv_s = nn.Conv2d(fin, fout, kernel_size=1, bias=False)# apply spectral norm if specifiedif 'spectral' in opt.norm_G:self.conv_0 = spectral_norm(self.conv_0)self.conv_1 = spectral_norm(self.conv_1)if self.learned_shortcut:self.conv_s = spectral_norm(self.conv_s)# define normalization layersspade_config_str = opt.norm_G.replace('spectral', '')self.norm_0 = SPADE(spade_config_str, fin, opt.semantic_nc)self.norm_1 = SPADE(spade_config_str, fmiddle, opt.semantic_nc)if self.learned_shortcut:self.norm_s = SPADE(spade_config_str, fin, opt.semantic_nc)# note the resnet block with SPADE also takes in |seg|,# the semantic segmentation map as inputdef forward(self, x, seg):x_s = self.shortcut(x, seg)dx = self.conv_0(self.actvn(self.norm_0(x, seg)))dx = self.conv_1(self.actvn(self.norm_1(dx, seg)))out = x_s + dxreturn outdef shortcut(self, x, seg):if self.learned_shortcut:x_s = self.conv_s(self.norm_s(x, seg))else:x_s = xreturn x_sdef actvn(self, x):return F.leaky_relu(x, 2e-1)

models/networks/normalization.py

class SPADE(nn.Module):def __init__(self, config_text, norm_nc, label_nc):super().__init__()assert config_text.startswith('spade')parsed = re.search('spade(\D+)(\d)x\d', config_text)param_free_norm_type = str(parsed.group(1))ks = int(parsed.group(2))if param_free_norm_type == 'instance':self.param_free_norm = nn.InstanceNorm2d(norm_nc, affine=False)elif param_free_norm_type == 'syncbatch':self.param_free_norm = SynchronizedBatchNorm2d(norm_nc, affine=False)elif param_free_norm_type == 'batch':self.param_free_norm = nn.BatchNorm2d(norm_nc, affine=False)else:raise ValueError('%s is not a recognized param-free norm type in SPADE'% param_free_norm_type)# The dimension of the intermediate embedding space. Yes, hardcoded.nhidden = 128pw = ks // 2self.mlp_shared = nn.Sequential(nn.Conv2d(label_nc, nhidden, kernel_size=ks, padding=pw),nn.ReLU())self.mlp_gamma = nn.Conv2d(nhidden, norm_nc, kernel_size=ks, padding=pw)self.mlp_beta = nn.Conv2d(nhidden, norm_nc, kernel_size=ks, padding=pw)def forward(self, x, segmap):# Part 1. generate parameter-free normalized activationsnormalized = self.param_free_norm(x)  # 与仿射变换参数无关的标准化# Part 2. produce scaling and bias conditioned on semantic mapsegmap = F.interpolate(segmap, size=x.size()[2:], mode='nearest')#对segmap做resizeactv = self.mlp_shared(segmap)gamma = self.mlp_gamma(actv)beta = self.mlp_beta(actv)# apply scale and biasout = normalized * (1 + gamma) + beta#这里解释一个为什么是1+gamma而不是gamma,作者自己解释是因为怕gamma学习到的结果接近于0,                #那乘以normalized以后就为0了,失去了normalized的作用,所以要用1+gamma,确保 #normalized有发挥作用同时还能学习仿射变换return out

这里我放一下paper里的网络图,可以对照代码看一下

(对应SPADE)

(左边对应SPADEResnetBlock,右边对应Generator)

最后,还是说一下我看这篇paper的一个疑惑问题,除了做语义图像合成,比如像上面这些代码,如果我们不用--use_vae的话,那训练的时候就是一个mask对应一个real image,最后学到的风格是一致的,还可以用来做不同风格的图像生成吗(考虑输入real image来影响生成结果)?看看作者在github的回复:

To produce outputs with different styles, you need to train with VAE by using --use_vae flag. It it was not trained with VAE, it cannot generate different styles.

The pretrained models of COCO, ADE20K and Cityscapes are all without VAE, because we actually didn't want random generation of styles, in order to keep the evaluation metric reproducible. As you know. for GauGAN video, we trained with VAE. Once you finish training with VAE, to produce different styles for the same semantic layout input, simply run the model multiple times. It will always generate different results.

如果你想要在同样的Mask上生成多种风格的结果,用--use_vae即可。

好吧,再看看use_vae做了什么

models/pix2pix_model.py

    def generate_fake(self, input_semantics, real_image, compute_kld_loss=False):z = NoneKLD_loss = Noneif self.opt.use_vae:z, mu, logvar = self.encode_z(real_image) ###在这里生成了zif compute_kld_loss:KLD_loss = self.KLDLoss(mu, logvar) * self.opt.lambda_kldfake_image = self.netG(input_semantics, z=z)assert (not compute_kld_loss) or self.opt.use_vae, \"You cannot compute KLD loss if opt.use_vae == False"return fake_image, KLD_loss
    def encode_z(self, real_image):mu, logvar = self.netE(real_image)z = self.reparameterize(mu, logvar)return z, mu, logvar
    def reparameterize(self, mu, logvar):std = torch.exp(0.5 * logvar)eps = torch.randn_like(std)return eps.mul(std) + mu

models/networks/encoder.py

class ConvEncoder(BaseNetwork):""" Same architecture as the image discriminator """def __init__(self, opt):super().__init__()kw = 3pw = int(np.ceil((kw - 1.0) / 2))ndf = opt.ngfnorm_layer = get_nonspade_norm_layer(opt, opt.norm_E)self.layer1 = norm_layer(nn.Conv2d(3, ndf, kw, stride=2, padding=pw))self.layer2 = norm_layer(nn.Conv2d(ndf * 1, ndf * 2, kw, stride=2, padding=pw))self.layer3 = norm_layer(nn.Conv2d(ndf * 2, ndf * 4, kw, stride=2, padding=pw))self.layer4 = norm_layer(nn.Conv2d(ndf * 4, ndf * 8, kw, stride=2, padding=pw))self.layer5 = norm_layer(nn.Conv2d(ndf * 8, ndf * 8, kw, stride=2, padding=pw))if opt.crop_size >= 256:self.layer6 = norm_layer(nn.Conv2d(ndf * 8, ndf * 8, kw, stride=2, padding=pw))self.so = s0 = 4self.fc_mu = nn.Linear(ndf * 8 * s0 * s0, 256)self.fc_var = nn.Linear(ndf * 8 * s0 * s0, 256)self.actvn = nn.LeakyReLU(0.2, False)self.opt = optdef forward(self, x):if x.size(2) != 256 or x.size(3) != 256:x = F.interpolate(x, size=(256, 256), mode='bilinear')x = self.layer1(x)x = self.layer2(self.actvn(x))x = self.layer3(self.actvn(x))x = self.layer4(self.actvn(x))x = self.layer5(self.actvn(x))if self.opt.crop_size >= 256:x = self.layer6(self.actvn(x))x = self.actvn(x)x = x.view(x.size(0), -1)mu = self.fc_mu(x)logvar = self.fc_var(x)return mu, logvar

encode具体就不分析。

写的比较乱,主要是为了梳理一下自己的思路,如果有错误还请评论指正。

这篇关于SPADE 代码略解 ade20k数据集的文章就介绍到这儿,希望我们推荐的文章对编程师们有所帮助!



http://www.chinasem.cn/article/781283

相关文章

springboot循环依赖问题案例代码及解决办法

《springboot循环依赖问题案例代码及解决办法》在SpringBoot中,如果两个或多个Bean之间存在循环依赖(即BeanA依赖BeanB,而BeanB又依赖BeanA),会导致Spring的... 目录1. 什么是循环依赖?2. 循环依赖的场景案例3. 解决循环依赖的常见方法方法 1:使用 @La

使用C#代码在PDF文档中添加、删除和替换图片

《使用C#代码在PDF文档中添加、删除和替换图片》在当今数字化文档处理场景中,动态操作PDF文档中的图像已成为企业级应用开发的核心需求之一,本文将介绍如何在.NET平台使用C#代码在PDF文档中添加、... 目录引言用C#添加图片到PDF文档用C#删除PDF文档中的图片用C#替换PDF文档中的图片引言在当

C#使用SQLite进行大数据量高效处理的代码示例

《C#使用SQLite进行大数据量高效处理的代码示例》在软件开发中,高效处理大数据量是一个常见且具有挑战性的任务,SQLite因其零配置、嵌入式、跨平台的特性,成为许多开发者的首选数据库,本文将深入探... 目录前言准备工作数据实体核心技术批量插入:从乌龟到猎豹的蜕变分页查询:加载百万数据异步处理:拒绝界面

用js控制视频播放进度基本示例代码

《用js控制视频播放进度基本示例代码》写前端的时候,很多的时候是需要支持要网页视频播放的功能,下面这篇文章主要给大家介绍了关于用js控制视频播放进度的相关资料,文中通过代码介绍的非常详细,需要的朋友可... 目录前言html部分:JavaScript部分:注意:总结前言在javascript中控制视频播放

Spring Boot 3.4.3 基于 Spring WebFlux 实现 SSE 功能(代码示例)

《SpringBoot3.4.3基于SpringWebFlux实现SSE功能(代码示例)》SpringBoot3.4.3结合SpringWebFlux实现SSE功能,为实时数据推送提供... 目录1. SSE 简介1.1 什么是 SSE?1.2 SSE 的优点1.3 适用场景2. Spring WebFlu

Java利用JSONPath操作JSON数据的技术指南

《Java利用JSONPath操作JSON数据的技术指南》JSONPath是一种强大的工具,用于查询和操作JSON数据,类似于SQL的语法,它为处理复杂的JSON数据结构提供了简单且高效... 目录1、简述2、什么是 jsONPath?3、Java 示例3.1 基本查询3.2 过滤查询3.3 递归搜索3.4

java之Objects.nonNull用法代码解读

《java之Objects.nonNull用法代码解读》:本文主要介绍java之Objects.nonNull用法代码,具有很好的参考价值,希望对大家有所帮助,如有错误或未考虑完全的地方,望不吝赐... 目录Java之Objects.nonwww.chinasem.cnNull用法代码Objects.nonN

SpringBoot实现MD5加盐算法的示例代码

《SpringBoot实现MD5加盐算法的示例代码》加盐算法是一种用于增强密码安全性的技术,本文主要介绍了SpringBoot实现MD5加盐算法的示例代码,文中通过示例代码介绍的非常详细,对大家的学习... 目录一、什么是加盐算法二、如何实现加盐算法2.1 加盐算法代码实现2.2 注册页面中进行密码加盐2.

MySQL大表数据的分区与分库分表的实现

《MySQL大表数据的分区与分库分表的实现》数据库的分区和分库分表是两种常用的技术方案,本文主要介绍了MySQL大表数据的分区与分库分表的实现,文中通过示例代码介绍的非常详细,对大家的学习或者工作具有... 目录1. mysql大表数据的分区1.1 什么是分区?1.2 分区的类型1.3 分区的优点1.4 分

Mysql删除几亿条数据表中的部分数据的方法实现

《Mysql删除几亿条数据表中的部分数据的方法实现》在MySQL中删除一个大表中的数据时,需要特别注意操作的性能和对系统的影响,本文主要介绍了Mysql删除几亿条数据表中的部分数据的方法实现,具有一定... 目录1、需求2、方案1. 使用 DELETE 语句分批删除2. 使用 INPLACE ALTER T