Linear regression using tf.data with federated core API and data on remote execution client

TFF does not support the mixing of federated computation with TensorFlow. The usual paradigm in TFF goes something like:

1. Write your local functions in TensorFlow, using TFF's @tff.tf_computation decorator
2. Inside the body of a tff.federated_computation, invoke intrinsic and communication operators (like tff.federated_broadcast above), and tff.federated_map your TF computations to federated elements.

There are many subtleties that have led to the design of a pattern like the above, but the main one is the desire for all logic to be expressible in a serialized representation, so that it can be split up and sent to the various devices in your federated system during true deployment. This is a little hand-wavey; you can think of TFF's decorators as defining a context in which the code you write will be traced, so that it can run in a totally platform-independent way.

Using this pattern, your computation above would look something like:

# These lines set up the TFF execution environment, and importantly are not used 
# during the function definitions below--only upon *invocation* of your function
# in the outer Python context will the execution environment be used.
channels = [grpc.insecure_channel(f'{ip_address}:{port}') for _ in range(10)]
tff.backends.native.set_remote_execution_context(channels, rpc_mode='STREAMING')


@tf.function()
def load_data():
    # From TFF's perspective, the code written here will be represented as a
    # "blob of Tensorflow" that can be shipped out anywhere. We don't actually
    # *need* a tf_computation decorator here, since this will be invoked in the
    # body of another tf_computation and the logic will get embedded there, but 
    # we could put one here if we wanted.
    return tf.data.experimental.CsvDataset(
        'data/Salary.csv', [tf.float64,tf.float64], header=True)

@tff.tf_computation(tf.float64, tf.float64)
@tf.function
def local_train(w, b):
    data = load_data()
    # We don't need a variable for the loss here, and though TFF will allow you
    # to construct a variable to use as a temporary in this fashion, tf.function
    # won't. We're pushing the TF team on that one ;).
    loss = tf.constant(0.0, dtype=tf.float64)
    with tf.GradientTape() as tape:
        # We must be inside a tf.function decorator, or the eager Python runtime,
        # to iterate over a tf.data.Dataset in this way.
        for x, y in data:
            p = x * w + b
            loss = loss + tf.square(p - y)

    g_w, g_b = tape.gradient(loss, [w, b])
    w = w- 0.0001 * g_w
    b = b- 0.0001 * g_b
    return w, b

# Making a symbol to represent these types is always a good idea
W_TYPE = tff.FederatedType(tf.float64, tff.CLIENTS, all_equal=True)
B_TYPE = tff.FederatedType(tf.float64, tff.CLIENTS, all_equal=True)

@tff.federated_computation(W_TYPE, B_TYPE)
def train(w, b):
    # Here w and b are elements of federated type, placed at clients.
    # We map the training function over these values.
    # If they were SERVER-placed instead, we would federated_broadcast them
    # out to the clients.
    updated_w, updated_b = tff.federated_map(local_train, (w, b))

# TFF's Python runtime represents federated values placed at clients as a list of
# values, with as many elements as there are clients. This number will be inferred 
# by the runtime and used to distribute the work. We also technically don't need
# variables here, but I believe they won't cause any problem.
clients_placed_w = [tf.Variable(2.0, dtype=tf.float64)]
clients_placed_b = [tf.Variable(3.0, dtype=tf.float64)]

# We could move the execution environment setup lines all the way down here,
# no problem.
for _ in range(1000):
    clients_placed_w, clients_placed_b = train(clients_placed_w, clients_placed_b)