Function run

examples/hfl/src/client.py:202–357

(
        global_rank,
        world_size,
        device_id,
        max_epoch,
        batch_size,
        model,
        data,
        data_dist,
        sgd,
        graph,
        verbosity,
        dist_option="plain",
        spars=None,
        precision="float32",
)

def run(
        global_rank,
        world_size,
        device_id,
        max_epoch,
        batch_size,
        model,
        data,
        data_dist,
        sgd,
        graph,
        verbosity,
        dist_option="plain",
        spars=None,
        precision="float32",
):
    # Connect to server
    client = Client(global_rank=device_id)
    client.start()

    dev = device.get_default_device()
    dev.SetRandSeed(0)
    np.random.seed(0)

    # Prepare dataset
    train_x, train_y, val_x, val_y, num_classes = get_data(data, data_dist, device_id)

    num_channels = train_x.shape[1]
    data_size = np.prod(train_x.shape[1: train_x.ndim]).item()

    # Prepare model
    model = get_model(
        model, num_channels=num_channels, num_classes=num_classes, data_size=data_size
    )

    if model.dimension == 4:
        image_size = train_x.shape[2]

    # For distributed training, sequential has better performance
    if hasattr(sgd, "communicator"):
        DIST = True
        sequential = True
    else:
        DIST = False
        sequential = False

    if DIST:
        train_x, train_y, val_x, val_y = partition(
            global_rank, world_size, train_x, train_y, val_x, val_y
        )

    if model.dimension == 4:
        tx = tensor.Tensor(
            (batch_size, num_channels, model.input_size, model.input_size),
            dev,
            singa_dtype[precision],
        )
    elif model.dimension == 2: