1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21

模型的分析更多见：https://zhuanlan.zhihu.com/p/159759699

### Profiler
用于分析训练性能的工具，需要用到 CUPTI，通过捕获一些 step 的各个算子计算时间来分析模型训练性能。  
发现的一个事实是： Batch Norm 的 FusedBatchNormV3 其实挺耗时间的，而且需要在 Host 做，不过换成其他 Normalization 方法其实也差不多。如果模型计算量本身比较小的话，那 GPU 利用率会降低不少。

## 方法
#### 计算
- `tf.reduce_sum()`：类似 `np.sum()`，区别在于它的输出保存在__相同的类型__的 Tensor 中，也就是说输入 `[254, 44, 44], dtype=uint8`，返回的值为 `86`，类型是 `uint8`，不会有溢出警告
- `tf.gather(data, idx)`：从 data 中取出 idx 号元素（整型 Tensor）形成一个 Tensor，也就是类似切片

#### 随机
为了让每次实验可重复，设置好随机是必须的，而设置确定的随机需要几个步骤：
1. 给参数初始化、数据集 on-the-fly 扩增等依赖随机的部分设置种子，并自动记录下来
2. 给 numpy、random 库设置全局种子
3. 在训练命令之前 `export TF_CUDNN_DETERMINISTIC=true` 可以让 cudnn 使用确定性算法（seed == 0）

#### 类型转换
```python
t2 = tf.cast(t1, dtype=tf.float32)

 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

## From numpy array, the first axis should be the same?
  # x_train: (60000, 28, 28, 1)
  # y_train: (60000,)
train_data = tf.data.Dataset.from_tensor_slices((x_train, y_train)).shuffle(60000).batch(32)
test_data = tf.data.Dataset.from_tensor_slices((x_test, y_test)).batch(32)

for images, labels in train_data:   # iterables
	pass

## From CSV
column_to_use = [0, 1, 2, 3, 4, 8]
record_defaults = [tf.int32, tf.int32, tf.string, tf.string, tf.float32, tf.float32]
data = tf.data.experimental.CsvDataset("titanic_dataset.csv", record_defaults, header=True, select_cols=column_to_use)



## From TFRecord
dataset = tf.data.TFRecordDataset(path)

  # Create a dictionary describing the features.
feature_desc = {
    'height': tf.io.FixedLenFeature([], tf.int64),
    'width': tf.io.FixedLenFeature([], tf.int64),
    'depth': tf.io.FixedLenFeature([], tf.int64),
    'image_raw': tf.io.FixedLenFeature([], tf.string),
    'mask_raw': tf.io.FixedLenFeature([], tf.string),
}

def _parse_image_function(example_proto):
  # Parse the input tf.Example proto using the dictionary above.
  return tf.io.parse_single_example(example_proto, feature_desc)

train = dataset.map(_parse_image_function)
# train['image_raw']: Tensor

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22

## Save/Load checkpoint
# Save in model fitting callback
cp_callback = tf.keras.callbacks.ModelCheckpoint(filepath=checkpoint_dir,
                                                 save_weights_only=True,
                                                 verbose=1)
# Manually save
model.save_weights(checkpoint_path.format(epoch=0))

model.load_weights(checkpoint_path)

## Save/Load entire model (resume training or ...)
# SavedModel format (default in TF2)
model.save('my_model') 

# keras h5 format
model.save('my_model.h5') 

tf.keras.models.load_model(xxx)    # compile with same arguments

model = create_model()
latest = tf.train.latest_checkpoint(checkpoint_dir)
model.load_weights(latest)   # load checkpoint

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22

model = tf.keras.Model(inputs=x, outputs=y)

## 1 Use Keras API to train
model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=0.001), 
              loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
              metrics=['accuracy'])
model.fit(train_dataset)


## 2 Use tf.GradientTape
@tf.function
def train_step(images, labels):
  with tf.GradientTape() as tape:
    # training=True is only needed if there are layers with different
    # behavior during training versus inference (e.g. Dropout).
    predictions = model(images, training=True)
    loss = loss_object(labels, predictions)
  gradients = tape.gradient(loss, model.trainable_variables)
  optimizer.apply_gradients(zip(gradients, model.trainable_variables))

  train_loss(loss)
  train_accuracy(labels, predictions)

 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

## Convert to TFLite
# Convert the model.
converter = tf.lite.TFLiteConverter.from_saved_model('efficientb0_UNet')
tflite_model = converter.convert()
# Save the TF Lite model.
with tf.io.gfile.GFile('efficientb0_UNet.tflite', 'wb') as f:
    f.write(tflite_model)
    
    
    
## Convert to Frozen Graph
from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2

# Convert Keras model to ConcreteFunction
full_model = tf.function(lambda x: model(x))
full_model = full_model.get_concrete_function(
    tf.TensorSpec(model.inputs[0].shape, model.inputs[0].dtype))

# Get frozen ConcreteFunction
frozen_func = convert_variables_to_constants_v2(full_model)
frozen_func.graph.as_graph_def()

layers = [op.name for op in frozen_func.graph.get_operations()]
print("-" * 50)
print("Frozen model layers: ")
for layer in layers:
    print(layer)

print("-" * 50)
print("Frozen model inputs: ")
print(frozen_func.inputs)
print("Frozen model outputs: ")
print(frozen_func.outputs)

# Save frozen graph from frozen ConcreteFunction to hard drive
tf.io.write_graph(graph_or_graph_def=frozen_func.graph,
                  logdir="./frozen_models",
                  name="frozen_graph.pb",
                  as_text=False)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16

psmodel = unet_model(2)

# For nested model
model.layers[1].save_weights('weights/efficientb0_UNet_81_f1score/efficient_sgd_81_30e_24bs_20200616_backbone.h5')
psmodel.layers[1].load_weights('weights/efficientb0_UNet_81_f1score/efficient_sgd_81_30e_24bs_20200616_backbone.h5')

# For each layer
for i in range(2, len(psmodel.layers)-1):
    if(type(psmodel.layers[i]) == tf.keras.Sequential):
        psmodel.layers[i].set_weights(model.layers[i].get_weights())

filter_81 = model.layers[-1].get_weights()[0][...,:2,:]
bias_81 = model.layers[-1].get_weights()[1][...,:2]
psmodel.layers[-1].set_weights([filter_81, bias_81])

psmodel.save('efficientb0_UNet_f1score/efficient_sgd_81to2_20200616-183119')

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23

## 以 metrics 方式，y_true 和 y_pred 都是 [batch_size, height, width, channel] 的 graph mode tensor
def f1score(y_true, y_pred):
    y_pred_ = tf.argmax(y_pred, axis=-1, output_type=tf.int64)[...,tf.newaxis]
    y_true_ = tf.cast(y_true, dtype=tf.int64)
    
    tmp_1 = tf.reduce_sum(tf.multiply(y_true_, y_pred_), axis=[1,2,3])
    tmp_2 = tf.reduce_sum(y_pred_, axis=[1,2,3])
    tmp_3 = tf.reduce_sum(y_true_, axis=[1,2,3])
    ttmp_1 = tf.cast(tmp_1, dtype=tf.float32)
    ttmp_2 = tf.cast(tmp_2, dtype=tf.float32)
    ttmp_3 = tf.cast(tmp_3, dtype=tf.float32)
    
    precision = ttmp_1 / (ttmp_2+0.000001)
    recall = ttmp_1 / (ttmp_3+0.000001)
    f1_score = 2. * precision * recall / (precision + recall + 0.000001)
    return f1_score

model.compile(..., loss=xxx, metrics=[f1score])

  # 读取时需要添加自定义对象
model = tf.keras.models.load_model('xxx', custom_objects={"f1score":f1score})

## 以 callback 方式

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18

## Rough
base_model = ...
base_model.trainable = False

# function api
def new_model():
	inputs = tf.keras.layers.Input(shape=[224, 224, 3])
	x = base_model(inputs)
	y = model_head(x)
    return tf.keras.Model(inputs=inputs, outputs=y)
model = new_model()

# or in sequential api
model = tf.keras.Sequential([
  base_model,
  global_average_layer,
  prediction_layer
])

1
2
3
4

if step == 120:
	tf.profiler.experimental.start(path)
if step == 180:
    tf.profiler.experimental.stop()

 1
 2
 3
 4
 5
 6
 7
 8
 9
10

insignificant = []
for i in range(len(model.layers)):
    for j in range(len(model.layers[i].weights)):
        if "kernel" in model.layers[i].weights[j].name:
            sum_k = np.mean(np.abs(model.layers[i].weights[j].numpy()), axis=(0,1,2))
            norm_sum_k = sum_k / sum_k.max()
            for v in range(norm_sum_k.shape[0]):
                if norm_sum_k[v] < 1e-3:
                    insignificant += [(i, j, sum_k[v], sum_k.max(), model.layers[i].weights[j].name)]
                    print(i, j, sum_k[v], sum_k.max(), model.layers[i].weights[j].name)

 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

# 单机多卡
mirrored_strategy = tf.distribute.MirroredStrategy()
with mirrored_strategy.scope():
  model = xxx
  model.compile(loss='mse', optimizer='sgd')

model.fit(dataset)

num_workers = 4


# 多机多卡
os.environ['TF_CONFIG'] = json.dumps({
    'cluster': {
        'worker': ["localhost:12345", "localhost:23456"]
    },
    'task': {'type': 'worker', 'index': 0}
})

strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy()
global_batch_size = per_worker_batch_size * num_workers
multi_worker_dataset = mnist_dataset(global_batch_size)

with strategy.scope():
  multi_worker_model = build_and_compile_cnn_model()

multi_worker_model.fit(multi_worker_dataset)

 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

strategy = tf.distribute.MirroredStrategy()
global_batch_size = per_worker_batch_size * strategy.num_replicas_in_sync

## 定义模型和优化器
with mirrored_strategy.scope():
    G = get_generator([flags.z_dim])
    D = get_discriminator([flags.output_size, flags.output_size, flags.c_dim])

    G.trainable=True
    D.trainable=True

    d_optimizer = tf.optimizers.Adam(flags.lr, beta_1=flags.beta1)
    g_optimizer = tf.optimizers.Adam(flags.lr, beta_1=flags.beta1)
    
## 分发数据集
with mirrored_strategy.scope():
    images, images_path = get_anime(flags.output_size, flags.n_epoch, GLOBAL_BATCH_SIZE)
    dist_images = mirrored_strategy.experimental_distribute_dataset(images)

## 定义损失函数，注意 reduce 需要用 global batch size
with mirrored_strategy.scope():
    def compute_loss(fake_logits, real_logits):
        real_d_l = tf.nn.sigmoid_cross_entropy_with_logits(tf.ones_like(real_logits), real_logits)
        fake_d_l = tf.nn.sigmoid_cross_entropy_with_logits(tf.zeros_like(fake_logits), fake_logits)

        d_loss = fake_d_l + real_d_l
        g_loss = tf.nn.sigmoid_cross_entropy_with_logits(tf.ones_like(fake_logits), fake_logits)

        # 将 loss 各个元素视为一个 example，累加起来除以 GLOBAL_BATCH_SIZE
        return tf.nn.compute_average_loss(d_loss, global_batch_size=GLOBAL_BATCH_SIZE), \
               tf.nn.compute_average_loss(g_loss, global_batch_size=GLOBAL_BATCH_SIZE)
               
## 定义训练函数
def step_fn(batch_images):
    with tf.GradientTape(persistent=True) as tape:
        z = tf.random.normal([flags.batch_size, flags.z_dim], mean=0.0, stddev=1.0)
        fake_images = G(z, training=True)

        fake_logits = D(fake_images, training=True)
        real_logits = D(batch_images, training=True)

        d_loss, g_loss = compute_loss(fake_logits, real_logits)

    g_grad = tape.gradient(g_loss, G.trainable_weights)
    g_optimizer.apply_gradients(zip(g_grad, G.trainable_weights))
    d_grad = tape.gradient(d_loss, D.trainable_weights)
    d_optimizer.apply_gradients(zip(d_grad, D.trainable_weights))
    del tape
    return d_loss, g_loss

@tf.function
def distributed_train_step(dataset_inputs):
    replica_d_losses, replica_g_losses = mirrored_strategy.run(step_fn, args=(dataset_inputs,))
    return mirrored_strategy.reduce(tf.distribute.ReduceOp.SUM, replica_d_losses, axis=None), \
           mirrored_strategy.reduce(tf.distribute.ReduceOp.SUM, replica_g_losses, axis=None)

## 训练
for epoch in range(flags.n_epoch):
    sum_d_loss = 0.0
    sum_g_loss = 0.0
    num_step = 0
    step_time = time.time()

    # 分发策略采用数据分发，每张卡有同等大小的 batch size
    for x in dist_images:
        # 去掉每个 epoch 最后一个不完全的 batch
        if x.values[0].shape[0] != BATCH_SIZE_PER_REPLICA:
            break
        tmp_d_loss, tmp_g_loss = distributed_train_step(x)
        sum_d_loss += tmp_d_loss
        sum_g_loss += tmp_g_loss
        num_step += 1

    d_loss = sum_d_loss / num_step
    g_loss = sum_g_loss / num_step

1
2
3

image_holder = tf.placeholder(tf.float32, [b, h, w, c])
label_holder = tf.placeholder(tf.int32, [batch_size])
keepprob_holder = tf.placeholder(tf.float32)

1

y_conv=inference_op(image_holder, keepprob_holder)

1

loss = loss(y_conv,label_holder)

1

train_step = tf.train.AdamOptimizer(1e-4).minimize(loss)

1

tf.global_variables_initializer().run()

1
2
3
4
5
6
7

for i in range(20000):
    # batch = mnist.train.next_batch(50)
    image_batch, label_batch = sess.run([images_train, labels_train])
    if i%10 == 0:
        train_accuracy = accuracy.eval(feed_dict={image_holder: image_batch, label_holder: label_batch, keepprob_holder: 1.0})
        print("step %d, training accuracy %g"%(i, train_accuracy))
    train_step.run(feed_dict={image_holder: image_batch, label_holder: label_batch, keepprob_holder: 0.5})

 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

from tensorflow.python.tools import freeze_graph 

gan = UGATIT(args)
gan.build_model()

# 可以写在 build_model() 里面
gan.test_domain_A = tf.placeholder(tf.float32, [1, gan.img_size, gan.img_size, gan.img_ch], name='test_domain_A')
gan.test_fake_B, _ = gan.generate_a2b(gan.test_domain_A)

# 指定 checkpoint
checkpoint_dir = os.path.join(gan.checkpoint_dir, gan.model_dir)
ckpt = tf.train.get_checkpoint_state(checkpoint_dir)
ckpt_name = os.path.basename(ckpt.model_checkpoint_path)
model_checkpoint_path = os.path.join(checkpoint_dir, ckpt_name)

with tf.Session() as sess:
    tf.train.write_graph(sess.graph_def, './models', 'MLP_deconv_avg.pb') 
    freeze_graph.freeze_graph(
        input_graph='./models/MLP_deconv_avg.pb',
        input_saver='',
        input_binary=False, 
        input_checkpoint=model_checkpoint_path, 
        output_node_names='generator_B/out',   # 输出的节点需要定义 name，最后 scope + name 为此处的值
        restore_op_name='save/restore_all',
        filename_tensor_name='save/Const:0',
        output_graph='./models/MLP_deconv_avg_frozen_model.pb',
        clear_devices=True,  # 是否清除训练时节点指定的设备
        initializer_nodes=''
        )