问题类型

Bug

来源

source

Tensorflow版本

tf 2.11

自定义代码

无

OS平台和发行版

Ubuntu 20.04

移动设备

• 无响应*

Python版本

3.8

Bazel版本

• 无响应*

GCC/编译器版本

• 无响应*

CUDA/cuDNN版本

11.2/8.1

GPU型号和内存大小

Nvidia GeForce RTX 2070 6233MB

当前行为？

RPC works fine with CPU but when switched to GPU, memory increases linearly.
(Memory stays constant after awhile with CPU)

When training model, memory will keep increasing until desktop freezes.

Code below to reproduce the issue on Google Colab.

独立代码重现问题

# import os
# os.environ["CUDA_VISIBLE_DEVICES"] = "-1"

import tensorflow as tf
import numpy as np
import portpicker
import matplotlib.pyplot as plt

import time
import os
import psutil
process = psutil.Process(os.getpid())

#tf.debugging.set_log_device_placement(True)

def test_simple():

    @tf.function(input_signature=[
        tf.TensorSpec([], tf.int32),
        tf.TensorSpec([], tf.int32)])
    @tf.autograph.experimental.do_not_convert  
    def remote_fn(a, b):
        return tf.add(a,b)

    port = portpicker.pick_unused_port()
    address = "localhost:{}".format(port)
    server = tf.distribute.experimental.rpc.Server.create("grpc", address)
    server.register("addition", remote_fn)
    server.start()

    client = tf.distribute.experimental.rpc.Client.create("grpc",
                address=address, name="test_client")

    a = np.ones(shape=[100, 256], dtype=np.int32)
    b = np.ones(shape=[100, 256], dtype=np.int32)

    hist = []
    for i in range(100000):
        result = client.call(
            args=[a,b],
            method_name="addition",
            output_specs=tf.TensorSpec((), tf.int32))
        if result.is_ok():
            result.get_value()
        # value = client.addition_blocking(a,b)
        # print(value)

        if i % 1000 == 0:
            print(i, ": ", process.memory_info().rss/1024/1024, "MB")
            hist.append(process.memory_info().rss/1024/1024)

    server.close()
    print("After shutdown: ", process.memory_info().rss/1024/1024)
    plt.plot(hist)
    plt.show()

test_simple()

相关日志输出

909000 :  1414.50390625 MB
910000 :  1414.5078125 MB
911000 :  1414.5078125 MB
912000 :  1414.5078125 MB
913000 :  1414.5078125 MB
914000 :  1414.5078125 MB
915000 :  1414.51171875 MB
916000 :  1414.51171875 MB
917000 :  1414.51171875 MB
918000 :  1414.515625 MB
919000 :  1415.6640625 MB
920000 :  1416.796875 MB
921000 :  1416.796875 MB
922000 :  1416.796875 MB
923000 :  1417.75 MB
924000 :  1418.7109375 MB
925000 :  1418.7109375 MB
926000 :  1418.7109375 MB
927000 :  1418.7109375 MB
928000 :  1418.71484375 MB
929000 :  1418.84765625 MB
930000 :  1418.84765625 MB
931000 :  1418.8515625 MB
932000 :  1418.85546875 MB
933000 :  1419.62890625 MB
934000 :  1419.87890625 MB
935000 :  1420.37890625 MB
936000 :  1420.53125 MB
937000 :  1420.65625 MB
938000 :  1421.83984375 MB
939000 :  1421.84375 MB
940000 :  1422.74609375 MB
941000 :  1422.84375 MB
942000 :  1422.84375 MB
943000 :  1423.3515625 MB
944000 :  1423.8984375 MB
945000 :  1423.91015625 MB
946000 :  1425.5078125 MB
947000 :  1426.109375 MB
948000 :  1426.34765625 MB
949000 :  1426.34765625 MB
950000 :  1426.3515625 MB
951000 :  1426.3515625 MB
952000 :  1427.05859375 MB
953000 :  1427.37890625 MB
954000 :  1427.3828125 MB
955000 :  1428.90234375 MB
956000 :  1429.69140625 MB
957000 :  1429.6953125 MB
958000 :  1430.3046875 MB
959000 :  1432.421875 MB
960000 :  1432.421875 MB
961000 :  1432.421875 MB
962000 :  1432.82421875 MB
963000 :  1432.82421875 MB
964000 :  1432.82421875 MB
965000 :  1433.1328125 MB
966000 :  1433.1328125 MB
967000 :  1433.1328125 MB
968000 :  1433.1328125 MB
969000 :  1434.60546875 MB
970000 :  1434.60546875 MB
971000 :  1434.640625 MB
972000 :  1435.08203125 MB
973000 :  1435.0859375 MB
974000 :  1435.59765625 MB
975000 :  1436.98828125 MB
976000 :  1436.9921875 MB
977000 :  1437.125 MB
978000 :  1437.98828125 MB
979000 :  1441.01953125 MB
980000 :  1441.96875 MB
981000 :  1441.96875 MB
982000 :  1441.97265625 MB
983000 :  1441.97265625 MB
984000 :  1441.97265625 MB
985000 :  1441.97265625 MB
986000 :  1441.97265625 MB
987000 :  1441.97265625 MB
988000 :  1441.97265625 MB
989000 :  1441.97265625 MB
990000 :  1441.9765625 MB
991000 :  1441.9765625 MB
992000 :  1441.98046875 MB
993000 :  1444.41796875 MB
994000 :  1444.78515625 MB
995000 :  1444.78515625 MB
996000 :  1444.78515625 MB
997000 :  1446.05078125 MB
998000 :  1446.05078125 MB
999000 :  1446.05078125 MB
After shutdown:  1447.05078125