PF-Net/model.py at master · Lifann/PF-Net · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
from context import ctx

import common
import snippet
import tensorflow as tf
import numpy as np


class Model(object):
  def __init__(self,
               x=tf.placeholder(tf.float32, shape=(ctx.num_incomplete, 3)),
               y_gt=tf.placeholder(tf.float32, shape=(ctx.num_cropped, 3))):
    """
    x: tensor. cropped input tensor.
    y_gt: tensor. cropped output tensor.
    """
    self.x = x
    self.y_gt = y_gt

    self.y_bold = None
    self.y_mid = None
    self.y_fine = None

    self.y_gt_bold = None
    self.y_gt_mid = None
    self.y_gt_fine = None

    self.g_loss = None
    self.ad_loss = None
    self.loss = None
    self.train_op = None

  def build(self):
    self._prepare_multi_resolution_inputs()
    feature_vec = snippet.MRE(self.x,
                              k=ctx.MRE_k,
                              nn_sizes=ctx.MRE_nn_sizes,
                              agg_num=ctx.MRE_agg_num)
    self.y_bold, self.y_mid, self.y_fine =      \
        snippet.PPD(feature_vec,
                    ctx.PPD_M,
                    M1=ctx.PPD_M1,
                    M2=ctx.PPD_M2,
                    FC_sizes=ctx.PPD_FC_sizes)

    self.g_loss = snippet.g_loss_fn(
        self.y_bold,
        self.y_mid,
        self.y_fine,
        self.y_gt_bold,
        self.y_gt_mid,
        self.y_gt_fine,
        eta=ctx.GLOSS_coef)

    self.ad_loss = snippet.ad_loss_fn(
        self.y_fine,
        self.y_gt_fine,
        CMLP_nn_sizes=ctx.ADLOSS_CMLP_nn_size,
        agg_num=ctx.ADLOSS_agg_num,
        nn_sizes=ctx.ADLOSS_nn_sizes)

    self.loss = ctx.loss_coef * self.g_loss  \
              + (1 - ctx.loss_coef) * self.ad_loss

    self.optimizer = tf.train.AdamOptimizer(learning_rate=ctx.learning_rate,
                                            beta1=ctx.beta1,
                                            beta2=ctx.beta2,
                                            epsilon=ctx.epsilon)

    self.train_op = self.optimizer.minimize(self.loss)

  def _prepare_multi_resolution_inputs(self):
    """
    Create multiple resolution inputs for y_gt.
    """
    self.y_gt_bold = tf.placeholder(tf.float32, shape=(ctx.PPD_M1, 3))
    self.y_gt_mid = tf.placeholder(tf.float32, shape=(ctx.PPD_M2, 3))
    self.y_gt_fine = self.y_gt