衣服生成项目

基于论文Semantically Multi-modal Image Synthesis，并且很详细

AI Studio

546人浏览 · 2022-02-11 22:31:09

AI Studio · 2022-02-11 22:31:09 发布

请点击此处查看本环境基本用法.

Please click here for more detailed instructions.

项目名称：衣服生成

（主要属于语义生成领域,GAN）

今天隆重推出我的衣服生成项目，鼓掌鼓掌。

基于论文Semantically Multi-modal Image Synthesis，这篇论文的讲解我已经详细的在我的这个项目进行了阐述，可谓相当详细了。

然后这个项目就是以上述这篇论文进行复现，并且基于具体任务进行了合理的修改，再给自己鼓个掌。
最后我也提供了测试代码，大家可以看看。

本项目的前世今生

我在这里稍微介绍一下这个项目到现在的前世今生，首先我拿到的是衣服生成项目这个任务的时候，我其实已经初步想法就是为设计行业赋能，更好的帮助设计师在设计的过程中就可以对于最终的设计进行一个合理的把控，当我逛B站看到了SPADE的展示效果的时候，这种马良画笔的感觉让我就立马感到很亲切，没错这个就是我的追求，王八看绿豆，看对眼了，形容这个时候的我很亲切，于是说干就干，我直接先找到SPADE的论文，先进行了一番阅读，参考它的pytorch代码，然后就直接写了一个paddle的。

在这个过程中我也看了FutureSI大佬的Spade复现，其中他个人对于gan的发展理解对我确实可以说是个非常不错的学习，但是这个时候我遇到了第一个问题，数据集。我最后选择了kaggle一个比赛的数据集(在数据集介绍时候有介绍)。

第二个问题就是发现训练很慢，1batch需要几十秒，这个时候我在思考，因为我跑的图片大小只是256*256，所以不太可能是前向传播tensor计算量的问题，这个时候问题来到了数据处理这里，这个时候我比较懒，就将就的训练起来了。结果发现spade效果不行，有可能迭代次数太少，这个时候我发现了这个项目主要依据的论文Semantically Multi-modal Image Synthesis，因为它用了deep fashion这个数据集并且效果很不错，然后也是基于spade，所以我基于spade修改其实挺简单的，于是就有了这个项目的雏形，在这过程中我开始第一次尝试离线resize，发现这样速度确实快很多，这里是GT的指导让我坚定了这个想法，然后确实很好用。另外离线resize保存语义分割信息的时候，我刚开始保存的太稀疏了[256，256，class_num]，这样很占内存。但是因为一个像素其实只有一个标签，所以最后我保存成[256,256,1]，这也是感谢GT指出了我这个小白错误。

另外我是spade用的是IN，这里遵循FutureSI的建议，可以提高一点点效果。

首先就两个字，展示：

It’s show time.
从左往右分别为模型生成效果，Ground Truth(可以理解为模型参考答案)，模型输入的语义分割可视化

任务要求详细阐述

首先介绍一下这个任务要求，得给大家讲解一下我想做的是一件什么事情，我想做的是衣服生成任务这件事。
但是衣服生成任务和其他任务区别之一就是要求输出的是只有一个干净的衣服，就相当于不要有花里胡哨的背景，也就是要求生成器要把注意力侧重放在衣服上，而不是生成出的一整个图片。

数据介绍及其处理过程：

我选取的数据是Kaggle的一个比赛FGVC6的数据，这里对于数据提供的比赛致以崇高的感谢，这个数据提供了精准的衣服各个部位的划分标注区域，可以直接看我上面这个效果展示的GT和语义分割效果可视化，这就是训练时输入模型的两部分数据。
测试时候模型输入就不需要这个GT了，只需要语义分割信息就够了，这个部分具体细节后面讲。
输入给模型语义分割Tensor格式为[batch_size,class_num,H,W],分别具体为[4,46,256,256]，46是具体有46个label,然后训练的时候batch_size为4。
另外就是我会在loss计算的时候考虑衣服的蒙版。
因为我输入的数据格式要求是256*256，因此我需要把图片和语义分割信息resize，让H和W都为256，原先比赛原先图片太大了H和W都是上千的，因此这个resize操作其实是很消耗时间的，这个时候就会有人会问了，大佬大佬为啥不用crop裁剪呢？在这里我说出几大原因：
1. 因为我们这个应用是衣服生成，所以模型最好考虑的是衣服这个整体，尽量少的给局部信息，防止管中窥豹，要让模型有格局。
2. 因为一张真实的图片中衣服占的位置并不那么大，位置并不固定，很大概率你在2500*2500上裁剪出一个256*256的区域大概率是全黑的，不提供信息。
3. 而且这样处理数据很容易导致一些标签模型可见机率很小，比如鞋子，因为这个鞋子所占比例很小。
这个时候，我用的是在线resize，这样其实训练很慢，这个数据预处理很费时间，然后我采取离线resize保存成npy这个技巧，这样就是说训练就马力全开。但是吧我保存的这个npy是[256,256,46]，很稀疏比较占内存，于是我就只保存了1000组左右，于是我把npy保存的向量保存成[256，256，1]，这样就好很多，最后为10000组左右。

import paddle
import os
import paddle.nn as nn 
import numpy as np 
import pandas as pd

## 这是从在线resize用的dataloader，结果发现不好用，所以这个代码块请一直尘封就行
# from build_dataloader_clothes import Get_paired_dataset
# data_loader = Get_paired_dataset(1)#直接从原始数据集读取数据

如果只是为了测试效果，其实不需要制作数据集，我提供了几组数据用于实验，直接跑最后一个代码块就行

# # resize时候的数据集存放地址：自己制作离线resize第一步，解压原始数据集,解压时间有点长
# if not os.path.isdir("./data/d"):
#     os.mkdir("./data/d")
# ! unzip data/data125914/imaterialist-fashion-2019-FGVC6.zip -d ./data/d

# # 自己制作离线resize第二步，离线resize npy保存代码，好了这离线就做完了，建议其实训练简单保存个1万组就行，
# from build_dataloader_clothes import pair_data
# dataset = pair_data()
# i = 0
# import paddle
# import numpy as np
# from tqdm import tqdm
# for input_img,mask in tqdm(dataset):
#     # print(input_img.shape,mask.shape)
#     mask = paddle.sum(paddle.arange(1,47).unsqueeze(0).unsqueeze(0) *paddle.ones([256,256,46])*paddle.to_tensor(mask),axis =2,keepdim =True).numpy()
#     # print(mask)
#     np.save("./all_dataset/imgs/"+str(i)+".npy",input_img)
#     np.save("./all_dataset/masks/"+str(i)+".npy",mask)
#     i += 1
#     # break

#测试VGG19，输出a,b a为最后一层特征图，b为列表，相当于把中间的特征图也输出了，也包括最后一层特征图，可详见我的公开项目  CV GAN模型常用方便套用
from VGG_Model import VGG19
import numpy as np
import paddle
m = np.random.random([1, 3,256,256])
real_image = paddle.to_tensor(m,dtype="float32")
VGG19()(real_image)

接下来类似这种 !python -u ENCODER.py就是单元测试,更好为了debug

# !python -u ENCODER.py

# !python -u SPADEResBlock.py

# !python -u MODEL.py

接下来讲解一下我的一些文件夹及其文件内容。

VGG_MODEL.py为VGG方便计算感知损失，VGG文件夹为保存VGG19需要的参数文件。
ENCODER.py 为模型训练的时候需要的编码器，Generator为模型主体解码器，然后模型解码器调用了SPADEResBlock.py的SPADEResBlock，然后SPADEResBlock调用了spade.py 的spade.
Normal.py主要为了谱归一化。
GANLOSS.py提供lsgan对抗损失构造，但是我没有，我用的是STgan,自己发明的东西，总体思想是判别器学习标准答案，生成器去向判别器学习。
Discriminator.py提供多尺寸判别器
all_dataset文件夹有两个文件夹分别保存mask和img的npy
训练时图片保存在kl_result文件夹中
MODEL中的Model为encoder和generator的封装，训练的时候直接用这个.
训练时候MODEL forward把 z = paddle.randn([1,46*8,8,8])注释掉，测试的时候取消注释表示丢掉encoder


# 这个就是使用离线的npy文件构造的dataset
import paddle 
import numpy as np
import os
class Dataset(paddle.io.Dataset):
    def __init__(self):
        self.root = "/home/aistudio/all_dataset"
        self.imgs_ori = os.path.join(self.root, "imgs")
        self.masks_ori = os.path.join(self.root, "masks")

        self.imgs_path_ori = Dataset.data_maker(self.imgs_ori)
        self.masks_path_ori = Dataset.data_maker(self.masks_ori)
        self.size = len(self.imgs_path_ori)

    @staticmethod
    def data_maker(dir):
        dir_list = []
        assert os.path.isdir(dir), '%s is not a valid directory' % dir

        for root, _, fnames in sorted(os.walk(dir)):
            for fname in fnames:
                if Dataset.is_npy_file(fname):
                    path = os.path.join(root, fname)
                    dir_list.append(path)

        return sorted(dir_list)

    @staticmethod    
    def is_npy_file(filename):
        return any(filename.endswith(extension) for extension in ["npy"])


    def __getitem__(self, index):
        input_img = np.load(self.imgs_path_ori[index])
        masks = np.load(self.masks_path_ori[index])
        masks = paddle.nn.functional.one_hot(paddle.to_tensor(masks).squeeze(2),47, name=None)[:,:,1:].numpy()#从[256,256,1]到[256,256,class_num]关键就在于one_hot函数

        return (input_img,masks)

    def __len__(self):
        return self.size

batch_size = 4
datas = Dataset()
data_loader =  paddle.io.DataLoader(datas,batch_size=batch_size,shuffle =True)
for input_img,masks in data_loader:
    print(input_img.shape,masks.shape)
    break

## 主要测试在线resize进行数据预处理消耗的时间。
# import time
# j = 0
# t = time.clock( ) 
# # for i in data_loader:
# #     print(j,"该轮耗时{:2f}s".format(time.clock()-t))
# #     t = time.clock()
# #     j+=1
# #     break

from ENCODER import ConvEncoder
from Generator import SPADEGenerator

import paddle
import paddle.nn as nn

KLDLoss主要实现这个：推导详见VAE的loss

编码器E产生了一个潜在的编码Z，该编码在训练过程中应该遵循一个高斯分布N（0,1）。在测试时，编码器E被丢弃。从高斯分布中随机抽样的编码代替z。为了实现这一点，我们使用重新参数化技巧[26]在训练过程中启用可微损失函数。具体来说，编码器通过两个全连接的层来预测一个平均向量和一个方差向量来表示编码的分布。编码z分布和高斯分布之间的差距可以通过施加kl-散度损失来最小化。


class KLDLoss(nn.Layer):
    def forward(self, mu, logvar):
        return -0.5 * paddle.sum(1 + logvar - mu.pow(2) - logvar.exp())
KLD_Loss = KLDLoss()
l1loss = nn.L1Loss()

# !python -u build_dataloader_clothes.py

from VGG_Model import VGG19
VGG = VGG19()

import paddle
import cv2
from tqdm import tqdm
import numpy as np
import os
from GANloss import GANLoss
from visualdl import LogWriter
from MODEL import Model
import math
log_writer = LogWriter("./log/gnet")
mse_loss = paddle.nn.MSELoss()
l1loss = paddle.nn.L1Loss()

# !python -u Discriminator.py
'''
该代码块代表多尺度判别器示例
'''
from Discriminator import build_m_discriminator
import numpy as np
discriminator = build_m_discriminator()
input_nc = 3
x = np.random.uniform(-1, 1, [4, 3, 256, 256]).astype('float32')
x = paddle.to_tensor(x)
print("input tensor x.shape",x.shape)\

y = discriminator(x)
for i in range(len(y)):
    for j in range(len(y[i])):
        print(i, j, y[i][j].shape)
    print('--------------------------------------')

encoder = ConvEncoder()
generator = SPADEGenerator()
model = Model()

#model和discriminator参数文件导入
# M_path ='model_params/Mmodel_state2.pdparams'
# layer_state_dictm = paddle.load(M_path)
# model.set_state_dict(layer_state_dictm)


# D_path ='discriminator_params/Dmodel_state2.pdparams'
# layer_state_dictD = paddle.load(D_path)
# discriminator.set_state_dict(layer_state_dictD)

scheduler_G = paddle.optimizer.lr.StepDecay(learning_rate=1e-4, step_size=3, gamma=0.8, verbose=True)
scheduler_D = paddle.optimizer.lr.StepDecay(learning_rate=4e-4, step_size=3, gamma=0.8, verbose=True)

optimizer_G = paddle.optimizer.Adam(learning_rate=scheduler_G,parameters=model.parameters(),beta1=0.,beta2 =0.9)
optimizer_D = paddle.optimizer.Adam(learning_rate=scheduler_D,parameters=discriminator.parameters(),beta1=0.,beta2 =0.9)

EPOCHEES = 100
i = 0

#四个设计保存参数文件的文件夹
save_dir_generator = "generator_params"
save_dir_encoder = "encoder_params"
save_dir_model = "model_params"
save_dir_Discriminator = "discriminator_params"

class Train_OPT():
    '''
    opt格式
    '''
    def __init__(self):
        super(Train_OPT, self).__init__()
        self.no_vgg_loss = False
        self.batchSize = 4
        self.lambda_feat = 10.0
        self.lambda_vgg = 2
opt = Train_OPT()

#单纯当个指标,实际style_loss不参与反向传播
def gram(x):
    b, c, h, w = x.shape
    x_tmp = x.reshape((b, c, (h * w)))
    gram = paddle.matmul(x_tmp, x_tmp, transpose_y=True)
    return gram / (c * h * w)

def style_loss(style, fake):
    mean_loss = paddle.sqrt(paddle.abs(paddle.square(paddle.mean(style))-paddle.square(paddle.mean(fake))))*0.5
    std_loss =  paddle.sqrt(paddle.abs(paddle.square(paddle.std(style))-paddle.square(paddle.std(fake))))*0.5

    gram_loss = nn.L1Loss()(gram(style), gram(fake))
    return gram_loss
    # return gram_loss

trans_img是为了实现Encoder的特殊输入，[b,3*46,256,256],具体原因详见我的项目论文解读一篇关于语义生成论文（要求控制单独语义生成）中的Enooder详解。

def trans_img(input_semantics, real_image):
    images = None
    seg_range = input_semantics.shape[1]
    for i in range(input_semantics.shape[0]):
        resize_image = None
        for n in range(0, seg_range):
            seg_image = real_image[i] * input_semantics[i][n]
            seg_image = seg_image.unsqueeze(axis=0)#[1,3,h,w]
                # resize_image = torch.cat((resize_image, seg_image), dim=0)
            if resize_image is None:
                resize_image = seg_image
            else:
                resize_image = paddle.concat((resize_image, seg_image), axis=1)
        if images is None:
            images = resize_image
        else:
            images = paddle.concat((images, resize_image), axis=0)
    return images

这里给大家简单讲解一下我的模型个人改进地方：

判别器有三个任务，它需要判别GT为True,判别生成器生成出的图片为False,判别语义分割可视化为False，这个就是我的改进了，这个是为了判别器帮助生成器生成更加复杂真实的纹理。
为了稍微侧重于有衣服的区域，最后生成器的featloss只考虑有衣服mask的部分。
然后我把spade.py中nn.conv2d(46,128)变成普通的卷积，没有使用分组卷积。因为46不能除尽group_num = 4

参与反向传播的所有g_loss部分：

g_loss = g_ganloss + g_vggloss +g_featloss +kldloss

再次强调：

训练时候MODEL forward把 z = paddle.randn([1,46*8,8,8])注释掉，测试的时候取消注释表示丢掉encoder

# 训练代码
step =0
for epoch in range(EPOCHEES):
    # if(step >1000):
        # break
    for input_img,mask in tqdm(data_loader):
        try:
            # if(step >1000):
            #     break
            # print(input_img.shape,mask.shape)
            input_img =paddle.transpose(x=input_img.astype("float32")/127.5-1,perm=[0,3,1,2])
            mask = paddle.transpose(x=mask,perm=[0,3,1,2]).astype("float32")
            seg_mask = (paddle.sum(paddle.arange(1,47).unsqueeze(0).unsqueeze(-1).unsqueeze(-1)*paddle.ones([opt.batchSize,46,256,256])*mask,axis =1,keepdim =True).astype("float32")*3+50)/255-1
            seg_mask = paddle.concat([seg_mask,seg_mask,seg_mask],axis =1)
            b,c,h,w = input_img.shape
            
            model_input = trans_img(mask,input_img)

            img_fake,_,_ = model(model_input,mask)
            img_fake = img_fake.detach()
            # kld_loss = KLD_Loss(mu,logvar)
            # print(img_fake.shape)

            fake_and_real_data = paddle.concat((img_fake, input_img,seg_mask), 0).detach()
            pred = discriminator(fake_and_real_data)

            df_ganloss = 0.
            for i in range(len(pred)):
                pred_i = pred[i][-1][:opt.batchSize]
                # new_loss = -paddle.minimum(-pred_i - 1, paddle.zeros_like(pred_i)).mean() # hingle loss pred_i<-1
                new_loss = (300 * 1.2 *GANLoss()(pred_i, False))/4
                df_ganloss += new_loss
            df_ganloss /= len(pred)
            df_ganloss*=0.35
            
            dr_ganloss = 0.
            for i in range(len(pred)):
                pred_i = pred[i][-1][opt.batchSize:opt.batchSize*2]
                # new_loss = -paddle.minimum(pred_i - 1, paddle.zeros_like(pred_i)).mean() # hingle loss  pred_i>1
                new_loss = (300 * 1.2 *GANLoss()(pred_i, True))/4
                dr_ganloss += new_loss
            dr_ganloss /= len(pred)
            dr_ganloss*=0.35

            dseg_ganloss = 0.
            for i in range(len(pred)):
                pred_i = pred[i][-1][opt.batchSize*2:]
                # new_loss = -paddle.minimum(pred_i - 1, paddle.zeros_like(pred_i)).mean() # hingle loss  pred_i>1
                new_loss = (300 * 1.2 *GANLoss()(pred_i, False))/4
                dseg_ganloss += new_loss
            dseg_ganloss /= len(pred)
            dseg_ganloss*=0.35

            d_loss = df_ganloss + dr_ganloss + dseg_ganloss


            d_loss.backward()
            optimizer_D.step()
            optimizer_D.clear_grad()

            discriminator.eval()
            # encoder.eval()
            # set_requires_grad(discriminator,False)
            # mu, logvar =  encoder(input_img)
            # kld_loss = KLD_Loss(mu,logvar)
            # z = reparameterize(mu, logvar)
            # img_fake = generator(mask,z)
            # print(img_fake.shape)
            img_fake,mu,logvar = model(model_input,mask)
            kldloss = KLD_Loss(mu,logvar)/20
            loss_mask = paddle.sum(mask,axis = 1,keepdim = True).astype("bool").astype("float32").detach()



            g_vggloss = paddle.to_tensor(0.)
            g_styleloss= paddle.to_tensor(0.)
            if not opt.no_vgg_loss:
                rates = [1.0 / 32, 1.0 / 16, 1.0 / 8, 1.0 / 4, 1.0]
                # _, fake_features = VGG( paddle.multiply (img_fake,loss_mask))
                # _, real_features = VGG(paddle.multiply (input_img,loss_mask))

                _, fake_features = VGG(img_fake)
                _, real_features = VGG(input_img)

                for i in range(len(fake_features)):
                    a,b = fake_features[i], real_features[i]
                    # if i ==len(fake_features)-1:
                    #     a = paddle.multiply( a,F.interpolate(loss_mask,a.shape[-2:]))
                    #     b = paddle.multiply( b,F.interpolate(loss_mask,b.shape[-2:]))
                    g_vggloss += rates[i] * l1loss(a,b)
                    # print(a.shape,b.shape)
                        # g_vggloss += paddle.mean(rates[i] *paddle.square(a-b))
                    if i ==len(fake_features)-1:
                        style_a,style_b = fake_features[i], real_features[i]
                        style_a = paddle.multiply( style_a,F.interpolate(loss_mask,style_a.shape[-2:]))
                        style_b = paddle.multiply( style_b,F.interpolate(loss_mask,style_b.shape[-2:]))
                        g_styleloss += rates[i] *  style_loss(style_b,style_a)
                    

                g_vggloss *= opt.lambda_vgg
                g_vggloss /=30

                g_styleloss/=50
            
            loss_mask8 = paddle.concat([loss_mask,loss_mask],axis=0)
            fake_and_real_data = paddle.concat((img_fake, input_img), 0)
            # fake_and_real_data = paddle.multiply (fake_and_real_data,loss_mask8)
            pred = discriminator(fake_and_real_data)
            # 关闭真图片 tensor 的梯度计算
            for i in range(len(pred)):
                for j in range(len(pred[i])):
                    pred[i][j][opt.batchSize:].stop_gradient = True

            g_ganloss = paddle.to_tensor(0.)
            for i in range(len(pred)):
                
                pred_i_f = pred[i][-1][:opt.batchSize]
                loss_mask0 = F.interpolate(loss_mask,pred_i_f.shape[-2:])
                # pred_i_f = paddle.multiply(pred_i_f,loss_mask0)

                pred_i_r = pred[i][-1][opt.batchSize:].detach()
                # pred_i_r = paddle.multiply(pred_i_r,loss_mask0)


                _,c,h,w = pred_i_f.shape
                # new_loss = -1*pred_i_f.mean() # hinge loss
                new_loss = paddle.sum(paddle.square(pred_i_r -pred_i_f))/math.sqrt(c*h*w)
                g_ganloss += new_loss
            g_ganloss /= len(pred)
            # g_ganloss*=20

            g_featloss = paddle.to_tensor(0.)
            for i in range(len(pred)):
                for j in range(len(pred[i]) - 1): # 除去最后一层的中间层featuremap
                    pred_i_f = pred[i][j][:opt.batchSize]
                    loss_mask0 = F.interpolate(loss_mask,pred_i_f.shape[-2:])
                    pred_i_f = paddle.multiply(pred_i_f,loss_mask0)

                    pred_i_r = pred[i][j][opt.batchSize:].detach()
                    pred_i_r = paddle.multiply(pred_i_r,loss_mask0)


                    unweighted_loss = (pred_i_r -pred_i_f).abs().mean() # L1 loss
                    g_featloss += unweighted_loss * opt.lambda_feat / len(pred)
            g_featloss*=3
    
            g_loss = g_ganloss  + g_vggloss +g_featloss +kldloss
            # g_loss = g_vggloss+g_styleloss
            g_loss.backward()
            optimizer_G.step()
            optimizer_G.clear_grad()

            # optimizer_E.step()
            # optimizer_E.clear_grad()        

            discriminator.train()

            if step%2==0:
                log_writer.add_scalar(tag='train/d_real_loss', step=step, value=dr_ganloss.numpy()[0])
                log_writer.add_scalar(tag='train/d_fake_loss', step=step, value=df_ganloss.numpy()[0])
                dseg_ganloss
                log_writer.add_scalar(tag='train/dseg_ganloss', step=step, value=dseg_ganloss.numpy()[0])
                log_writer.add_scalar(tag='train/d_all_loss', step=step, value=d_loss.numpy()[0])


                log_writer.add_scalar(tag='train/g_ganloss', step=step, value=g_ganloss.numpy()[0])
                log_writer.add_scalar(tag='train/g_featloss', step=step, value=g_featloss.numpy()[0])
                log_writer.add_scalar(tag='train/g_vggloss', step=step, value=g_vggloss.numpy()[0])
                log_writer.add_scalar(tag='train/g_loss', step=step, value=g_loss.numpy()[0])
                log_writer.add_scalar(tag='train/g_styleloss', step=step, value=g_styleloss.numpy()[0])
                log_writer.add_scalar(tag='train/kldloss', step=step, value=kldloss.numpy()[0])



            step+=1
            # print(i)
            if step%300 == 3:
                print(step,"g_ganloss",g_ganloss.numpy()[0],"g_featloss",g_featloss.numpy()[0],"g_vggloss",g_vggloss.numpy()[0],"g_styleloss",g_styleloss.numpy()[0],"kldloss",kldloss.numpy()[0],"g_loss",g_loss.numpy()[0])
                print(step,"dreal_loss",dr_ganloss.numpy()[0],"dfake_loss",df_ganloss.numpy()[0],"dseg_ganloss",dseg_ganloss.numpy()[0],"d_all_loss",d_loss.numpy()[0])

                # img_fake = paddle.multiply (img_fake,loss_mask)
                input_img = paddle.multiply (input_img,loss_mask)

                g_output = paddle.concat([img_fake,input_img,seg_mask],axis = 3).detach().numpy()                      # tensor -> numpy
                g_output = g_output.transpose(0, 2, 3, 1)[0]             # NCHW -> NHWC
                g_output = (g_output+1) *127.5                        # 反归一化
                g_output = g_output.astype(np.uint8)
                cv2.imwrite(os.path.join("./kl_result", 'epoch'+str(step).zfill(3)+'.png'),cv2.cvtColor(g_output,cv2.COLOR_RGB2BGR))
                # generator.train()
            
            if step%100 == 3:
                save_param_path_g = os.path.join(save_dir_generator, 'Gmodel_state'+str(3)+'.pdparams')
                paddle.save(model.generator.state_dict(), save_param_path_g)
                save_param_path_d = os.path.join(save_dir_Discriminator, 'Dmodel_state'+str(3)+'.pdparams')
                paddle.save(discriminator.state_dict(), save_param_path_d)
                # save_param_path_e = os.path.join(save_dir_encoder, 'Emodel_state'+str(1)+'.pdparams')
                # paddle.save(model.encoder.state_dict(), save_param_path_e)
                save_param_path_m = os.path.join(save_dir_model, 'Mmodel_state'+str(3)+'.pdparams')
                paddle.save(model.state_dict(), save_param_path_m)
            # break
        except:
            pass
        # break
    scheduler_G.step()
    scheduler_D.step()

loss可视化

#测试代码 效果保存至test文件
from MODEL import Model
import paddle
import numpy as np
import cv2
import os

def trans_img(input_semantics, real_image):
    images = None
    seg_range = input_semantics.shape[1]
    for i in range(input_semantics.shape[0]):
        resize_image = None
        for n in range(0, seg_range):
            seg_image = real_image[i] * input_semantics[i][n]
            seg_image = seg_image.unsqueeze(axis=0)#[1,3,h,w]
                # resize_image = torch.cat((resize_image, seg_image), dim=0)
            if resize_image is None:
                resize_image = seg_image
            else:
                resize_image = paddle.concat((resize_image, seg_image), axis=1)
        if images is None:
            images = resize_image
        else:
            images = paddle.concat((images, resize_image), axis=0)
    return images

model = Model(1)
M_path ='Mmodel_state3 (1).pdparams'
layer_state_dictm = paddle.load(M_path)
model.set_state_dict(layer_state_dictm)
input_img =paddle.to_tensor( np.load("all_dataset/imgs/1.npy")).astype("float32").unsqueeze(0)
# print(input_img.shape)
mask = np.load("all_dataset/masks/1.npy")
mask = paddle.nn.functional.one_hot(paddle.to_tensor(mask).squeeze(2),47, name=None)[:,:,1:].astype("float32").unsqueeze(0)

input_img =paddle.transpose(x=input_img.astype("float32")/127.5-1,perm=[0,3,1,2])
mask = paddle.transpose(x=mask,perm=[0,3,1,2]).astype("float32")
loss_mask = paddle.sum(mask,axis = 1,keepdim = True).astype("bool").astype("float32").detach()

seg_mask = paddle.sum(paddle.arange(1,47).unsqueeze(0).unsqueeze(-1).unsqueeze(-1)*paddle.ones([1,46,256,256])*mask,axis =1,keepdim =True).astype("float32")*5/255-1
seg_mask = paddle.concat([seg_mask,seg_mask,seg_mask],axis =1)

model_input = trans_img(mask,input_img)

img_fake,_,_ = model(model_input,mask)
print('img_fake',img_fake.shape)
input_img = paddle.multiply (input_img,loss_mask)
print(img_fake.shape,input_img.shape,seg_mask.shape)
g_output = paddle.concat([img_fake,input_img,seg_mask],axis = 3).detach().numpy()                      # tensor -> numpy
g_output = g_output.transpose(0, 2, 3, 1)[0]             # NCHW -> NHWC
g_output = (g_output+1) *127.5                        # 反归一化
g_output = g_output.astype(np.uint8)
cv2.imwrite(os.path.join("./test", str(1008)+'.png'), cv2.cvtColor(g_output,cv2.COLOR_RGB2BGR))