NVlabs/noise2noise代码（三）网络训练代码解析

蔚落 2022-04-23 13:14 741阅读 0赞

**目的：**看懂网络训练代码，方便后续更改。

**目录**

一、更改迭代次数

1.1定义位置

train\_config初始化位置

EasyDict定义位置

1.2 更改迭代次数的方法

二、网络结构

2.1 原始网络结构与代码解析

2.2 训练时autoencoder的调用

三、训练函数嵌套关系

3.1 config.py 到submit\_run

3.2 submit.py 中的run\_wrapper

3.3 调用train.py

3.4打印出信息

四、训练函数

4.1 输入参数及类型

4.2 初始化

4.3 创建网络

4.4 运用梯度更新及迭代

--------------------

# 一、更改迭代次数 #

原始迭代次数较大，所以运行一次耗时较久，不易调试，我们将迭代次数改小以便调试。直接在main下面的def train(args)中更改就可以。

## 1.1定义位置 ##

config.py之中，这里将其值改小就可以取得减少程序运行时间的效果。

if __name__ == "__main__":
        def train(args):
            if args:
                n2n = args.noise2noise if 'noise2noise' in args else True
                train_config.noise2noise = n2n
                if 'long_train' in args and args.long_train:
                    #train_config.iteration_count = 500000
                    train_config.iteration_count = 500
                    #train_config.eval_interval = 5000
                    train_config.eval_interval = 50
                    train_config.ramp_down_perc = 0.5

### train\_config初始化位置 ###

train_config = dnnlib.EasyDict(
        iteration_count=300000,
        eval_interval=1000,
        minibatch_size=4,
        run_func_name="train.train",
        learning_rate=0.0003,
        ramp_down_perc=0.3,
        noise=gaussian_noise_config,
    #    noise=poisson_noise_config,
        noise2noise=True,
        train_tfrecords='datasets/imagenet_val_raw.tfrecords',
        validation_config=default_validation_config
    )

可以看出将其定义为dnnlib.EasyDict格式的数据。类似于字典，但比字典更易使用。

### EasyDict定义位置 ###

train\_config与validate\_config都运用了dnnlib.EasyDict作为初始化，此class的位置为dnnlib/util.py

class EasyDict(dict):
        """Convenience class that behaves like a dict but allows access with the attribute syntax."""
    
        def __getattr__(self, name: str) -> Any:
            try:
                return self[name]
            except KeyError:
                raise AttributeError(name)
    
        def __setattr__(self, name: str, value: Any) -> None:
            self[name] = value
    
        def __delattr__(self, name: str) -> None:
            del self[name]

## 1.2 更改迭代次数的方法 ##

迭代次数分别改为,评估与输出的interval，和输出的斜率。

train_config.iteration_count = 100
                    train_config.eval_interval = 10
                    train_config.ramp_down_perc = 0.5

迭代次数设为500时，最终PSNR=27.34，时间为2m28s

迭代次数为100时，最终PSNR=23.42，时间为1m48s

# 二、网络结构 #

## 2.1 原始网络结构与代码解析 ##

原始网络可以看作一个unet

skips = [x]
    
        n = x
        n = conv_lr('enc_conv0', n, 48)
        n = conv_lr('enc_conv1', n, 48)
        n = maxpool2d(n)
        skips.append(n)
    
        n = conv_lr('enc_conv2', n, 48)
        n = maxpool2d(n)
        skips.append(n)
    
        n = conv_lr('enc_conv3', n, 48)
        n = maxpool2d(n)
        skips.append(n)
    
        n = conv_lr('enc_conv4', n, 48)
        n = maxpool2d(n)
        skips.append(n)
    
        n = conv_lr('enc_conv5', n, 48)
        n = maxpool2d(n)
        n = conv_lr('enc_conv6', n, 48)
    
        #-----------------------------------------------
        n = upscale2d(n)
        n = tf.concat([n, skips.pop()], axis=1)
        n = conv_lr('dec_conv5', n, 96)
        n = conv_lr('dec_conv5b', n, 96)
    
        n = upscale2d(n)
        n = tf.concat([n, skips.pop()], axis=1)
        n = conv_lr('dec_conv4', n, 96)
        n = conv_lr('dec_conv4b', n, 96)
    
        n = upscale2d(n)
        n = tf.concat([n, skips.pop()], axis=1)
        n = conv_lr('dec_conv3', n, 96)
        n = conv_lr('dec_conv3b', n, 96)
    
        n = upscale2d(n)
        n = tf.concat([n, skips.pop()], axis=1)
        n = conv_lr('dec_conv2', n, 96)
        n = conv_lr('dec_conv2b', n, 96)
    
        n = upscale2d(n)
        n = tf.concat([n, skips.pop()], axis=1)
        n = conv_lr('dec_conv1a', n, 64)
        n = conv_lr('dec_conv1b', n, 32)
    
        n = conv('dec_conv1', n, 3, gain=1.0)

结构如下，相应数值和尺寸需要按照上面代码进行变化。

![watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3dlaXhpbl8zNjQ3NDgwOQ_size_16_color_FFFFFF_t_70][]

基本可以确定autoencoder之中确定的网络结构。

## 2.2 训练时autoencoder的调用 ##

直接在autoencoder之中加输出的信息

(n2n) jcx@smart-dsp:~/Desktop/NVlabs_noise2noise$ CUDA_VISIBLE_DEVICES=0 python config.py --desc='-test' train --train-tfrecords=datasets/part_bsd300.tfrecords --long-train=false --noise=gaussian
    ----------train in config.py
    ----------Iteration count is 100 and eval_interval is 10
    {'iteration_count': 100, 'eval_interval': 10, 'minibatch_size': 4, 'run_func_name': 'train.train', 'learning_rate': 0.0003, 'ramp_down_perc': 0.5, 'noise': {'func_name': 'train.AugmentGaussian', 'train_stddev_rng_range': (0.0, 50.0), 'validation_stddev': 25.0}, 'noise2noise': True, 'train_tfrecords': 'datasets/part_bsd300.tfrecords', 'validation_config': {'dataset_dir': 'datasets/kodak'}}
    ----------submit_run in submit.py
    ----------submit_config.submit_target in {SubmitTarget.LOCAL},create new dir to run
    Creating the run dir: results/00010-autoencoder-test-n2n
    ----------_populate_run_dir function in submit.py. Copying files to the run dir
    ----------run_wrapper function in submit.py
    dnnlib: Running train.train() on localhost...
    ----------train in train.py
    ----------Setting up dataset source from datasets/part_bsd300.tfrecords
    ----------net = tflib.Network(**config.net_config) in train.py
    -------------autoencoder in network.py
    
    autoencoder                 Params      OutputShape             WeightShape
    ---                         ---         ---                     ---
    x                           -           (?, 3, 256, 256)        -
    enc_conv0                   1344        (?, 48, 256, 256)       (3, 3, 3, 48)
    enc_conv1                   20784       (?, 48, 256, 256)       (3, 3, 48, 48)
    MaxPool                     -           (?, 48, 128, 128)       -
    enc_conv2                   20784       (?, 48, 128, 128)       (3, 3, 48, 48)
    MaxPool_1                   -           (?, 48, 64, 64)         -
    enc_conv3                   20784       (?, 48, 64, 64)         (3, 3, 48, 48)
    MaxPool_2                   -           (?, 48, 32, 32)         -
    enc_conv4                   20784       (?, 48, 32, 32)         (3, 3, 48, 48)
    MaxPool_3                   -           (?, 48, 16, 16)         -
    enc_conv5                   20784       (?, 48, 16, 16)         (3, 3, 48, 48)
    MaxPool_4                   -           (?, 48, 8, 8)           -
    enc_conv6                   20784       (?, 48, 8, 8)           (3, 3, 48, 48)
    Upscale2D                   -           (?, 48, 16, 16)         -
    dec_conv5                   83040       (?, 96, 16, 16)         (3, 3, 96, 96)
    dec_conv5b                  83040       (?, 96, 16, 16)         (3, 3, 96, 96)
    Upscale2D_1                 -           (?, 96, 32, 32)         -
    dec_conv4                   124512      (?, 96, 32, 32)         (3, 3, 144, 96)
    dec_conv4b                  83040       (?, 96, 32, 32)         (3, 3, 96, 96)
    Upscale2D_2                 -           (?, 96, 64, 64)         -
    dec_conv3                   124512      (?, 96, 64, 64)         (3, 3, 144, 96)
    dec_conv3b                  83040       (?, 96, 64, 64)         (3, 3, 96, 96)
    Upscale2D_3                 -           (?, 96, 128, 128)       -
    dec_conv2                   124512      (?, 96, 128, 128)       (3, 3, 144, 96)
    dec_conv2b                  83040       (?, 96, 128, 128)       (3, 3, 96, 96)
    Upscale2D_4                 -           (?, 96, 256, 256)       -
    dec_conv1a                  57088       (?, 64, 256, 256)       (3, 3, 99, 64)
    dec_conv1b                  18464       (?, 32, 256, 256)       (3, 3, 64, 32)
    dec_conv1                   867         (?, 3, 256, 256)        (3, 3, 32, 3)
    ---                         ---         ---                     ---
    Total                       991203
    
    ----------Building TensorFlow graph...
    -------------autoencoder in network.py
    ----------train_step = opt.apply_updates() in train.py
    ----------Training...
    -------------autoencoder in network.py
    -------------autoencoder in network.py
    -------------autoencoder in network.py
    Average PSNR: 7.15
    iter 0          time 6s           sec/eval 0.0     sec/iter 0.00    maintenance 6.5
    Average PSNR: 18.83
    iter 10         time 37s          sec/eval 1.7     sec/iter 0.17    maintenance 28.3
    Average PSNR: 19.80
    iter 20         time 45s          sec/eval 1.1     sec/iter 0.11    maintenance 7.4
    Average PSNR: 21.02
    iter 30         time 53s          sec/eval 1.0     sec/iter 0.10    maintenance 7.2
    Average PSNR: 21.56
    iter 40         time 1m 01s       sec/eval 1.1     sec/iter 0.11    maintenance 7.1
    Average PSNR: 22.35
    iter 50         time 1m 09s       sec/eval 1.1     sec/iter 0.11    maintenance 6.9
    Average PSNR: 21.74
    iter 60         time 1m 18s       sec/eval 1.1     sec/iter 0.11    maintenance 7.5
    Average PSNR: 22.93
    iter 70         time 1m 26s       sec/eval 1.1     sec/iter 0.11    maintenance 6.6
    Average PSNR: 22.64
    iter 80         time 1m 34s       sec/eval 1.0     sec/iter 0.10    maintenance 7.7
    Average PSNR: 23.02
    iter 90         time 1m 43s       sec/eval 1.0     sec/iter 0.10    maintenance 7.7
    Elapsed time: 1m 52s
    dnnlib: Finished train.train() in 1m 53s.
    ----------try in run wrapper finally handle _finished.txt

# 三、训练函数嵌套关系 #

![watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3dlaXhpbl8zNjQ3NDgwOQ_size_16_color_FFFFFF_t_70 1][]

## 3.1 config.py 到submit\_run ##

config.py之中是主函数，训练过程中，函数的作用是加入相应参数之后，运行train函数

submit_config.run_desc = desc + args.desc
        if args.run_dir_root is not None:
            submit_config.run_dir_root = args.run_dir_root
        if args.command is not None:
            args.func(args)
        else:
            # Train if no subcommand was given
            train(args)

train函数中关键语句为

dnnlib.submission.submit.submit_run(submit_config, **train_config)

其中输入参数，train\_config与submit\_config就是上面设置的相应参数.其中，train\_config是一系列与训练相关的代码，submit\_config在submit.py之中，用于实现将运行目录及结果写入等等操作。

#train config
    train_config = dnnlib.EasyDict(
        iteration_count=300000,
        eval_interval=1000,
        minibatch_size=4,
        run_func_name="train.train",
        learning_rate=0.0003,
        ramp_down_perc=0.3,
        noise=gaussian_noise_config,
    #    noise=poisson_noise_config,
        noise2noise=True,
        train_tfrecords='datasets/imagenet_val_raw.tfrecords',
        validation_config=default_validation_config
    )
    
    #submit config
    submit_config = dnnlib.SubmitConfig()
    submit_config.run_dir_root = 'results'
    submit_config.run_dir_ignore += ['datasets', 'results']
    
    desc = "autoencoder"

## 3.2 submit.py 中的run\_wrapper ##

submit\_run之中,创建相应文件夹，在文件夹之中运行该运行的函数，核心语句为run\_wrapper函数

run\_wrapper函数中用于运行封装好的函数

核心语句：

try:
            print("dnnlib: Running {0}() on {1}...".format(submit_config.run_func_name, submit_config.host_name))
            start_time = time.time()
            util.call_func_by_name(func_name=submit_config.run_func_name, submit_config=submit_config, **submit_config.run_func_kwargs)
            print("dnnlib: Finished {0}() in {1}.".format(submit_config.run_func_name, util.format_time(time.time() - start_time)))
        except:
            if is_local:
                raise
            else:
                traceback.print_exc()
    
                log_src = os.path.join(submit_config.run_dir, "log.txt")
                log_dst = os.path.join(get_path_from_template(submit_config.run_dir_root), "{0}-error.txt".format(submit_config.run_name))
                shutil.copyfile(log_src, log_dst)
        finally:
            open(os.path.join(submit_config.run_dir, "_finished.txt"), "w").close()

其中，就是运用call\_func\_by\_name函数进行函数的调用与实现

util.call_func_by_name(func_name=submit_config.run_func_name, submit_config=submit_config, **submit_config.run_func_kwargs)

## 3.3 调用train.py ##

util.call\_func\_by\_name(func\_name=submit\_config.run\_func\_name, submit\_config=submit\_config, \*\*submit\_config.run\_func\_kwargs)

通过call\_func\_by\_name这个函数调用相应的函数，

def call_func_by_name(*args, func_name: str = None, **kwargs) -> Any:
        """Finds the python object with the given name and calls it as a function."""
        assert func_name is not None
        func_obj = get_obj_by_name(func_name)
        assert callable(func_obj)
        return func_obj(*args, **kwargs)

## 3.4打印出信息 ##

通过打印出信息验证了嵌套及调用关系

----------train in config.py
    ----------Iteration count is 100 and eval_interval is 10
    {'iteration_count': 100, 'eval_interval': 10, 'minibatch_size': 4, 'run_func_name': 'train.train', 'learning_rate': 0.0003, 'ramp_down_perc': 0.5, 'noise': {'func_name': 'train.AugmentGaussian', 'train_stddev_rng_range': (0.0, 50.0), 'validation_stddev': 25.0}, 'noise2noise': True, 'train_tfrecords': 'datasets/part_bsd300.tfrecords', 'validation_config': {'dataset_dir': 'datasets/kodak'}}
    ----------submit_run in submit.py
    ----------submit_config.submit_target in {SubmitTarget.LOCAL},create new dir to run
    Creating the run dir: results/00503-autoencoder-n2n
    ----------_populate_run_dir function in submit.py. Copying files to the run dir
    ----------run_wrapper function in submit.py
    dnnlib: Running train.train() on localhost...
    ---------------train function in train.py

# 四、训练函数 #

train.py用于进行函数的训练

![watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3dlaXhpbl8zNjQ3NDgwOQ_size_16_color_FFFFFF_t_70 2][]

## 4.1 输入参数及类型 ##

def train(
        submit_config: dnnlib.SubmitConfig,
        iteration_count: int,
        eval_interval: int,
        minibatch_size: int,
        learning_rate: float,
        ramp_down_perc: float,
        noise: dict,
        validation_config: dict,
        train_tfrecords: str,
        noise2noise: bool):

输入网络训练，验证以及batch size与iteration\_count等等参数，后面此参数会传入各种train中的函数，以用于定义网络训练。

## 4.2 初始化 ##

![20190311163949199.png][]

#create validation set
        noise_augmenter = dnnlib.util.call_func_by_name(**noise)
        validation_set = ValidationSet(submit_config)
        validation_set.load(**validation_config)
        print("----------train in train.py")
    
        # Create a run context (hides low level details, exposes simple API to manage the run)
        ctx = dnnlib.RunContext(submit_config, config)
    
        # Initialize TensorFlow graph and session using good default settings
        tfutil.init_tf(config.tf_config)
    
        #create dataset
        dataset_iter = create_dataset(train_tfrecords, minibatch_size, noise_augmenter.add_train_noise_tf)

## 4.3 创建网络 ##

![20190311164119680.png][]

# Construct the network using the Network helper class and a function defined in config.net_config
        with tf.device("/gpu:0"):
            print("----------net = tflib.Network(**config.net_config) in train.py")
            net = tflib.Network(**config.net_config)
    
        # Optionally print layer information
        net.print_layers()
    
        print('----------Building TensorFlow graph...')
        with tf.name_scope('Inputs'), tf.device("/cpu:0"):
            lrate_in        = tf.placeholder(tf.float32, name='lrate_in', shape=[])
    
            noisy_input, noisy_target, clean_target = dataset_iter.get_next()
            noisy_input_split = tf.split(noisy_input, submit_config.num_gpus)
            noisy_target_split = tf.split(noisy_target, submit_config.num_gpus)
            clean_target_split = tf.split(clean_target, submit_config.num_gpus)
    
        # Define the loss function using the Optimizer helper class, this will take care of multi GPU
        opt = tflib.Optimizer(learning_rate=lrate_in, **config.optimizer_config)
    
        #compute meansq_error
        for gpu in range(submit_config.num_gpus):
            with tf.device("/gpu:%d" % gpu):
                net_gpu = net if gpu == 0 else net.clone()
    
                denoised = net_gpu.get_output_for(noisy_input_split[gpu])
    
                if noise2noise:
                    meansq_error = tf.reduce_mean(tf.square(noisy_target_split[gpu] - denoised))
                else:
                    meansq_error = tf.reduce_mean(tf.square(clean_target_split[gpu] - denoised))
                # Create an autosummary that will average over all GPUs
                with tf.control_dependencies([autosummary("Loss", meansq_error)]):
                    opt.register_gradients(meansq_error, net_gpu.trainables)

## 4.4 运用梯度更新及迭代 ##

![watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3dlaXhpbl8zNjQ3NDgwOQ_size_16_color_FFFFFF_t_70 3][]

#apply updates
        print("----------train_step = opt.apply_updates() in train.py")
        train_step = opt.apply_updates()
    
        # Create a log file for Tensorboard
        summary_log = tf.summary.FileWriter(submit_config.run_dir)
        summary_log.add_graph(tf.get_default_graph())
    
        print('----------Training...')
        time_maintenance = ctx.get_time_since_last_update()
        ctx.update(loss='run %d' % submit_config.run_id, cur_epoch=0, max_epoch=iteration_count)
    
        # The actual training loop
        for i in range(iteration_count):
            # Whether to stop the training or not should be asked from the context
            if ctx.should_stop():
                break
    
            # Dump training status
            if i % eval_interval == 0:
    
                time_train = ctx.get_time_since_last_update()
                time_total = ctx.get_time_since_start()
    
                # Evaluate 'x' to draw a batch of inputs
                [source_mb, target_mb] = tfutil.run([noisy_input, clean_target])
                denoised = net.run(source_mb)
                #save_image(submit_config, denoised[0], "img_{0}_y_pred.png".format(i))
                #save_image(submit_config, target_mb[0], "img_{0}_y.png".format(i))
                #save_image(submit_config, source_mb[0], "img_{0}_x_aug.png".format(i))
    
                validation_set.evaluate(net, i, noise_augmenter.add_validation_noise_np)
    
                print('iter %-10d time %-12s sec/eval %-7.1f sec/iter %-7.2f maintenance %-6.1f' % (
                    autosummary('Timing/iter', i),
                    dnnlib.util.format_time(autosummary('Timing/total_sec', time_total)),
                    autosummary('Timing/sec_per_eval', time_train),
                    autosummary('Timing/sec_per_iter', time_train / eval_interval),
                    autosummary('Timing/maintenance_sec', time_maintenance)))
    
                dnnlib.tflib.autosummary.save_summaries(summary_log, i)
                ctx.update(loss='run %d' % submit_config.run_id, cur_epoch=i, max_epoch=iteration_count)
                time_maintenance = ctx.get_last_update_interval() - time_train
    
            lrate =  compute_ramped_down_lrate(i, iteration_count, ramp_down_perc, learning_rate)
            tfutil.run([train_step], {lrate_in: lrate})

[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3dlaXhpbl8zNjQ3NDgwOQ_size_16_color_FFFFFF_t_70]: /images/20220317/895c0a6b43b5482b83affb4994973479.png
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3dlaXhpbl8zNjQ3NDgwOQ_size_16_color_FFFFFF_t_70 1]: /images/20220317/0859a9c2069c4eb9b2778726631ccfd2.png
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3dlaXhpbl8zNjQ3NDgwOQ_size_16_color_FFFFFF_t_70 2]: /images/20220317/14cdd39bc4eb4472bc07161709b3d298.png
[20190311163949199.png]: /images/20220317/6b81b0f4efc3445d81a9316cddf71e7b.png
[20190311164119680.png]: /images/20220317/eab4a4340f9f4a6b81835ff137036873.png
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3dlaXhpbl8zNjQ3NDgwOQ_size_16_color_FFFFFF_t_70 3]: /images/20220317/0343ab7dab9943819dc625d65035d1c3.png