self.conv1=nn.Conv2d(3, 8 * 64, 3, padding=1)

    out=self.conv1(z)
    out_h=F.interpolate(out, [16, 16])

import torch
from torch.autograd import Variable # torch 中 Variable 模块

从checkpoint.pt中生成图像

import argparse
import functools
import os
from os import path
import json
import math
from collections import defaultdict
import random

import numpy as np
import torch
import torch.optim as optim
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader

from sg2im.data import imagenet_deprocess_batch
# from sg2im.data.coco import CocoSceneGraphDataset, coco_collate_fn
from sg2im.data.vg import VgSceneGraphDataset, vg_collate_fn
from sg2im.discriminators import PatchDiscriminator, AcCropDiscriminator
from sg2im.losses import get_gan_losses
from sg2im.metrics import jaccard
from sg2im.model import Sg2ImModel
from sg2im.utils import int_tuple, float_tuple, str_tuple
from sg2im.utils import timeit, bool_flag, LossManager
from imageio import imwrite

pathdir='/1T/ysh/sg2im-master/models/vg128output_with_model.pt'
checkpoint=torch.load(pathdir)

train_samples = checkpoint['train_samples'] #list 100 4 32*3*64*64
val_samples = checkpoint['val_samples']


tdimen = np.array(train_samples).shape[0]
print('train batch num',tdimen)

vdimen = np.array(val_samples).shape[0]
print('val batch nun',vdimen)

# Save the generated images
output_dir='/1T/ysh/sg2im-master/outputs/vg128output/train'

# for i in range(train_samples.shape[0]): #100
for i in range(1): #100
    gt_img = train_samples[i]['gt_img']
    gt_box_gt_mask = train_samples[i]['gt_box_gt_mask']
    gt_img_dir=os.path.join(output_dir, 'gt_img')
    gbm_dir=os.path.join(output_dir, 'gt_box_gt_mask')
    if not path.exists(gt_img_dir):
        os.makedirs(gt_img_dir)
    if not path.exists(gbm_dir):
        os.makedirs(gbm_dir)
    for j in range(gt_img.shape[0]): #batch size 32
        gt_imgs = gt_img[j].numpy().transpose(1, 2, 0)
        # 需要格式为（H,W,C）
        img_path1 = os.path.join(gt_img_dir, 'img%d_%d.png' % (i,j))

        imwrite(img_path1, gt_imgs)

        gt_box_gt_masks = gt_box_gt_mask[j].numpy().transpose(1, 2, 0)
        # 需要格式为（H,W,C）
        img_path2 = os.path.join(gbm_dir, 'img%d_%d.png' % (i,j))

        imwrite(img_path2, gt_box_gt_masks)

import argparse
import os, sys
from os import path
import time
import copy
import torch
from torch import nn
import numpy as np
import random
from torchsummary import summary
import shutil
import scipy.io as sio

import warnings
warnings.filterwarnings("ignore")

def seed_torch(seed=1029):
    random.seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)  # if you are using multi-GPU.
    torch.backends.cudnn.benchmark = False
    torch.backends.cudnn.deterministic = True


seed_torch(999)


from gan_training import utils
from gan_training.train import Trainer, update_average
from gan_training.toggle_ImageNet import toggle_grad_G, toggle_grad_D
from gan_training.logger import Logger
from gan_training.checkpoints import CheckpointIO
from gan_training.inputs import get_dataset
from gan_training.distributions import get_ydist, get_zdist
from gan_training.eval import Evaluator
from gan_training.config import (
    load_config, build_models, build_optimizers, build_lr_scheduler, build_models_PRE,
)


def get_parameter_number(net):
    total_num = sum(p.numel() for p in net.parameters())
    trainable_num = sum(p.numel() for p in net.parameters() if p.requires_grad)
    print('Total=', total_num, 'Trainable=', trainable_num, 'fixed=', total_num-trainable_num)


def load_part_model(m_fix, m_ini):
    dict_fix = m_fix.state_dic()
    dict_ini = m_ini.state_dic()

    dict_fix = {k: v for k, v in dict_fix.items() if k in dict_ini and k.find('embedding')==-1 and k.find('fc') == -1}
    dict_ini.update(dict_fix)
    m_ini.load_state_dict(dict_ini)
    return m_ini


def model_equal_all(model, dict):
    model_dict = model.state_dict()
    model_dict.update(dict)
    model.load_state_dict(model_dict)
    return model


def model_equal_part(model, dict_all):
    model_dict = model.state_dict()
    dict_fix = {k: v for k, v in dict_all.items() if k in model_dict and k.find('embedding') == -1 and k.find('fc') == -1}
    model_dict.update(dict_fix)
    model.load_state_dict(model_dict)
    return model


''' ===================--- Set the traning mode ---==========================
DATA: going to train
DATA_FIX: used as a fixed pre-trained model
G_Layer_FIX, D_Layer_FIX: number of layers to fix
============================================================================='''

DATA = 'Flowers'
image_path = './data/102flowers/all8189images/'
is_control_kernel = True

DATA_FIX = 'ImageNet'
Num_epoch = 3

main_path = './'
load_dir = './pretrained_model/'

if is_control_kernel:
    out_path = main_path+ DATA + '_our_AdaFM/'
else:
    out_path = main_path + DATA + '_not_AdaFM/'

config_path = main_path+'/configs/' +DATA+ '.yaml'


for choose in range(1):

    G_Layer_FIX = -4
    D_Layer_FIX = 2

    config = load_config(config_path, 'configs/default.yaml')

    config['generator']['layers'] = G_Layer_FIX
    config['discriminator']['layers'] = D_Layer_FIX
    config['data']['train_dir'] = image_path
    config['data']['test_dir'] = image_path

    config['generator']['name'] = 'resnet2_AdaFM'
    config['discriminator']['name'] = 'resnet2_AdaFM'

    config['training']['out_dir'] = out_path + 'G_%d_D_%d/'%(-G_Layer_FIX, D_Layer_FIX)
    if not os.path.isdir(config['training']['out_dir']):
        os.makedirs(config['training']['out_dir'])

    if 1:
        # Short hands
        batch_size = config['training']['batch_size']
        d_steps = config['training']['d_steps']
        restart_every = config['training']['restart_every']
        inception_every = config['training']['inception_every']
        save_every = config['training']['save_every']
        backup_every = config['training']['backup_every']
        sample_nlabels = config['training']['sample_nlabels']

        # out_dir = config['training']['out_dir']
        out_dir = '/1T/ysh/GANTransferLimitedData-master/output'
        checkpoint_dir = path.join(out_dir, 'chkpts')

        # Create missing directories
        if not path.exists(out_dir):
            os.makedirs(out_dir)
        if not path.exists(checkpoint_dir):
            os.makedirs(checkpoint_dir)
        shutil.copyfile(sys.argv[0], out_dir + '/training_script.py')

        # Logger
        checkpoint_io = CheckpointIO(
            checkpoint_dir=checkpoint_dir
        )

        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

        # Dataset
        train_dataset, nlabels = get_dataset(
            name=config['data']['type'],
            data_dir=config['data']['train_dir'],
            size=config['data']['img_size'],
            lsun_categories=config['data']['lsun_categories_train']
        )
        train_loader = torch.utils.data.DataLoader(
                train_dataset,
                batch_size=batch_size,
                num_workers=config['training']['nworkers'],
                shuffle=True, pin_memory=True, sampler=None, drop_last=True
        )

        # Number of labels
        nlabels = min(nlabels, config['data']['nlabels'])
        sample_nlabels = min(nlabels, sample_nlabels)





        # Create models
        ''' --------- Choose the fixed layer ---------------'''
        generator, discriminator = build_models(config)

        dict_G = torch.load(load_dir + DATA_FIX + 'Pre_generator')
        generator = model_equal_part(generator, dict_G)
        dict_D = torch.load(load_dir + DATA_FIX + 'Pre_discriminator')
        discriminator = model_equal_part(discriminator, dict_D)

        for name, param in generator.named_parameters():
            if name.find('small') >= 0:
                param.requires_grad = True
            else:
                param.requires_grad = False
            if name.find('small_adafm_') >= 0:
                param.requires_grad = False
        get_parameter_number(generator)
        
        for param in discriminator.parameters():
            param.requires_grad = False

        #toggle_grad_G(generator, True, G_Layer_FIX)
        toggle_grad_D(discriminator, True, D_Layer_FIX)





        # Put models on gpu if needed
        generator, discriminator = generator.to(device), discriminator.to(device)
        g_optimizer, d_optimizer = build_optimizers(generator, discriminator, config)

        # summary(generator, input_size=[(256,), (1,)])
        # summary(discriminator, input_size=[(3, 128, 128), (1,)])

        # Register modules to checkpoint
        checkpoint_io.register_modules(
            generator=generator,
            discriminator=discriminator,
            g_optimizer=g_optimizer,
            d_optimizer=d_optimizer,
        )

        # Logger
        logger = Logger(
            log_dir=path.join(out_dir, 'logs'),
            img_dir=path.join(out_dir, 'imgs'),
            monitoring=config['training']['monitoring'],
            monitoring_dir=path.join(out_dir, 'monitoring')
        )

        # Distributions
        ydist = get_ydist(nlabels, device=device)
        zdist = get_zdist(config['z_dist']['type'], config['z_dist']['dim'],
                          device=device)

        # Save for tests
        ntest = 100
        x_real, ytest = utils.get_nsamples(train_loader, ntest)
        ytest.clamp_(None, nlabels-1)
        ytest = ytest.to(device)
        ztest = zdist.sample((ntest,)).to(device)
        utils.save_images(x_real, path.join(out_dir, 'real.png'))
        print('saved %d samples'%ntest)

        # Test generator
        if config['training']['take_model_average']:
            generator_test = copy.deepcopy(generator)
            checkpoint_io.register_modules(generator_test=generator_test)
        else:
            generator_test = generator

        # Evaluator

        NNN = 8000
        x_real, _ = utils.get_nsamples(train_loader, NNN)
        evaluator = Evaluator(generator_test, zdist, ydist,
                              batch_size=batch_size, device=device,
                              fid_real_samples=x_real, inception_nsamples=NNN, fid_sample_size=NNN)
        # Train
        tstart = t0 = time.time()


        it = -1
        epoch_idx = -1

        # Reinitialize model average if needed
        if (config['training']['take_model_average']
                and config['training']['model_average_reinit']):
            update_average(generator_test, generator, 0.)

        # Learning rate anneling
        g_scheduler = build_lr_scheduler(g_optimizer, config, last_epoch=it)
        d_scheduler = build_lr_scheduler(d_optimizer, config, last_epoch=it)

        # Trainer
        trainer = Trainer(
            generator, discriminator, g_optimizer, d_optimizer,
            gan_type=config['training']['gan_type'],
            reg_type=config['training']['reg_type'],
            reg_param=config['training']['reg_param'],
            D_fix_layer=config['discriminator']['layers']
        )
        

    # Training loop
    print('Start training...')
    save_dir = config['training']['out_dir'] + '/models/'
    print('save_dir：',save_dir)
    
    if not os.path.isdir(save_dir):
        os.makedirs(save_dir)
    FLAG = 500

    inception_mean_all = []
    inception_std_all = []
    fid_all = []
    
	#数据集
	pathdir='/1T/ysh/sg2im-master/models/vg128output_with_model.pt'
	checkpoint=torch.load(pathdir)
	
	train_samples = checkpoint['train_samples'] #list 100 4 32*3*64*64
	#val_samples = checkpoint['val_samples']
	tdimen = np.array(train_samples).shape[0]
	print('train samples batch num',tdimen)

    for epoch_idx in range(Num_epoch):
        # epoch_idx += 1
        print('Start epoch %d...' % epoch_idx)

        for i in range(tdimen):
            gt_img = train_samples[i]['gt_img']
   			gt_box_gt_mask = train_samples[i]['gt_box_gt_mask']
        	
            it += 1
            g_scheduler.step()
            d_scheduler.step()

            d_lr = d_optimizer.param_groups[0]['lr']
            g_lr = g_optimizer.param_groups[0]['lr']
            # logger.add('learning_rates', 'discriminator', d_lr, it=it)
            # logger.add('learning_rates', 'generator', g_lr, it=it)

            x_real, z = gt_img.to(device), gt_box_gt_mask.to(device)
            y=torch.full([batch_size],0)
            #y.clamp_(None, nlabels-1)

            # Generators updates
            # z = zdist.sample((batch_size,))
            
            gloss, x_fake = trainer.generator_trainstep(y, z, FLAG + 1.0)
            FLAG = FLAG * 0.9995

            if config['training']['take_model_average']:
                update_average(generator_test, generator,
                               beta=config['training']['model_average_beta'])

            # Discriminator updates
            dloss, reg = trainer.discriminator_trainstep(x_real, y, x_fake)

            if is_control_kernel:
                if it == 10000:
                    for name, param in generator.named_parameters():
                        if name.find('small_adafm_') >= 0:
                            param.requires_grad = True
                    get_parameter_number(generator)

            with torch.no_grad():
                # (i) Sample if necessary
                if (it % config['training']['sample_every']) == 0:
                    d_fix, d_update = discriminator.conv_img.weight[1, 1, 1, 1], discriminator.fc.weight[0, 1]
                    g_fix, g_update = generator.conv_img.weight[1, 1, 1, 1], 0.0

                    print('[epoch %0d, it %4d] g_loss = %.4f, d_loss = %.4f, reg=%.4f, d_fix=%.4f, d_update=%.4f, g_fix=%.4f, g_update=%.4f'
                          % (epoch_idx, it, gloss, dloss, reg, d_fix, d_update, g_fix, g_update))
                    # print('Creating samples...')
                    x, _ = evaluator.create_samples(ztest, ytest)
                    logger.add_imgs(x, 'all', it, nrow=10)

                # (ii) Compute inception if necessary
                if inception_every > 0 and ((it + 2) % inception_every) == 0:
                    inception_mean, inception_std, fid = evaluator.compute_inception_score()
                    inception_mean_all.append(inception_mean)
                    inception_std_all.append(inception_std)
                    fid_all.append(fid)
                    print('test it %d: IS: mean %.2f, std %.2f, FID: mean %.2f, time: %2f' % (
                        it, inception_mean, inception_std, fid, time.time() - tstart))

                    FID = np.stack(fid_all)
                    Inception_mean = np.stack(inception_mean_all)
                    Inception_std = np.stack(inception_std_all)
                    sio.savemat(config['training']['out_dir'] + DATA + 'base_FID_IS.mat', {'FID': FID,
                                                           'Inception_mean': Inception_mean,
                                                           'Inception_std': Inception_std})

                # (iii) Backup if necessary
                if ((it + 1) % backup_every) == 0:
                    print('Saving backup...')
                    TrainModeSave = DATA + '_%08d_' % it
                    torch.save(generator_test.state_dict(), save_dir + TrainModeSave + 'Pre_generator')
                    torch.save(discriminator.state_dict(), save_dir + TrainModeSave + 'Pre_discriminator')