Cấu trúc đơn giản

#!/usr/bin/env python3
# -*- coding: UTF-8 -*-
import numpy as np
import tensorflow as tf

# Định nghĩa mô hình mạng nơ-ron
def build_network(input_data):
    layer = input_data
    layer = tf.layers.dense(layer, units=32, activation=tf.nn.relu)
    output = layer
    return output

# Các tham số siêu mô hình
batch_size = 20000
learning_rate = 0.001
num_epochs = 1000

# Định nghĩa đầu vào và nhãn
input_placeholder = tf.placeholder(tf.float32, shape=[None, 784])
target_placeholder = tf.placeholder(tf.float32, shape=[None])

# Tính toán hàm mất mát
predictions = build_network(input_placeholder)
cross_entropy_loss = tf.nn.softmax_cross_entropy_with_logits(
    labels=target_placeholder,
    logits=predictions
)

# Tối ưu hóa
optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate)
training_step = optimizer.minimize(cross_entropy_loss)

# Khởi tạo biến
init_op = tf.global_variables_initializer()

# Ghi log đồ thị
summary_writer = tf.summary.FileWriter(logdir='logs')

# Thực thi trong phiên làm việc
with tf.Session() as session:
    session.run(init_op)
    summary_writer.add_graph(graph=session.graph)
    summary_writer.flush()

    # Dữ liệu giả lập
    sample_data = np.random.rand(100, 11)

    steps_per_epoch = int(sample_data.shape[0] / batch_size)

    for epoch in range(num_epochs):
        for step in range(steps_per_epoch):
            start_idx = step * batch_size
            end_idx = (step + 1) * batch_size
            batch_x = sample_data[start_idx:end_idx, 0].reshape(-1, 784)
            batch_y = sample_data[start_idx:end_idx, 1:]

            _, loss_value = session.run(
                [training_step, cross_entropy_loss],
                feed_dict={
                    input_placeholder: batch_x,
                    target_placeholder: batch_y
                }
            )
            print(f"Epoch {epoch:5d}, Step {step:5d}: Loss = {loss_value}")

Sử dụng nhiều GPU hiệu quả

#!/usr/bin/env python3
# -*- coding: UTF-8 -*-
import numpy as np
import tensorflow as tf
from tensorflow.python.client import device_lib

# Đếm số GPU khả dụng
def detect_gpu_count():
    devices = device_lib.list_local_devices()
    gpu_list = [d.name for d in devices if d.device_type == 'GPU']
    count = len(gpu_list)
    print(f"{count} GPU(s) detected: {gpu_list}")
    return count

# Mô hình mạng với hỗ trợ tái sử dụng biến
def create_model(data_input, dropout_rate=1.0, reuse_vars=False):
    x = data_input
    with tf.variable_scope('hidden_layer', reuse=reuse_vars):
        x = tf.layers.dense(x, units=32, activation=tf.nn.relu)
    return x

# Cấu hình học tập
batch_size = 20000
learning_rate = 0.001
max_epochs = 1000
num_gpus = detect_gpu_count()

# Đầu vào dữ liệu
input_tensor = tf.placeholder(tf.float32, shape=[None, 784])
label_tensor = tf.placeholder(tf.float32, shape=[None])

# Chia dữ liệu theo số GPU
split_inputs = tf.split(input_tensor, num_gpus)
split_labels = tf.split(label_tensor, num_gpus)

# Tính toán mất mát trên từng GPU
total_losses = []
for i in range(num_gpus):
    with tf.device(f'/device:GPU:{i}'):
        with tf.variable_scope(tf.get_variable_scope(), reuse=(i > 0)):
            pred = create_model(split_inputs[i], reuse_vars=(i > 0))
            loss = tf.nn.softmax_cross_entropy_with_logits(
                labels=split_labels[i],
                logits=pred
            )
            total_losses.append(loss)

# Tổng hợp mất mát từ tất cả GPU
combined_loss = tf.reduce_mean(tf.concat(total_losses, axis=0))

# Tối ưu hóa toàn bộ
opt = tf.train.AdamOptimizer(learning_rate=learning_rate)
train_operation = opt.minimize(combined_loss)

# Khởi tạo biến
init_op = tf.global_variables_initializer()

# Ghi log đồ thị
writer = tf.summary.FileWriter(logdir='logs')

# Chạy mô hình
with tf.Session() as sess:
    sess.run(init_op)
    writer.add_graph(sess.graph)
    writer.flush()

    mock_data = np.random.rand(100, 11)
    steps_per_epoch = int(mock_data.shape[0] / batch_size)

    for epoch in range(max_epochs):
        for step in range(steps_per_epoch):
            start = step * batch_size
            end = (step + 1) * batch_size
            x_batch = mock_data[start:end, 0].reshape(-1, 784)
            y_batch = mock_data[start:end, 1:]

            _, loss_val = sess.run(
                [train_operation, combined_loss],
                feed_dict={
                    input_tensor: x_batch,
                    label_tensor: y_batch
                }
            )
            print(f"Epoch {epoch:5d}, Step {step:5d}: Loss = {loss_val}")

Tham khảo:

TensorFlow-Multi-GPUs/many-GPUs-MNIST.py tại master · golbin/TensorFlow-Multi-GPUs https://github.com/golbin/TensorFlow-Multi-GPUs/blob/master/many-GPUs-MNIST.py