from util.utils import send_msg, recv_msg, get_indices_each_node_case
# Configurations are in a separate config.py file
from config import *
model = get_model(model_name)
if hasattr(model, 'create_graph'):
model.create_graph(learning_rate=step_size)
if time_gen is not None:
use_fixed_averaging_slots = True
else:
use_fixed_averaging_slots = False
if batch_size < total_data: # Read all data once when using stochastic gradient descent
train_image, train_label, test_image, test_label, train_label_orig = get_data(dataset, total_data, dataset_file_path)
# This function takes a long time to complete,
# putting it outside of the sim loop because there is no randomness in the current way of computing the indices
indices_each_node_case = get_indices_each_node_case(n_nodes, MAX_CASE, train_label_orig)
listening_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
listening_sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
listening_sock.bind((SERVER_ADDR, SERVER_PORT))
client_sock_all=[]
# Establish connections to each client, up to n_nodes clients
while len(client_sock_all) < n_nodes:
listening_sock.listen(5)
print("Waiting for incoming connections...")
(client_sock, (ip, port)) = listening_sock.accept()
model_name
) # Used for computing loss_w_prev_min_loss for stochastic gradient descent,
# so that the state of model can be still used by control algorithm later.
if hasattr(model, 'create_graph'):
model.create_graph(learning_rate=step_size)
if hasattr(model2, 'create_graph'):
model2.create_graph(learning_rate=step_size)
# Assume the dataset does not change
if read_all_data_for_stochastic or batch_size >= total_data:
if batch_size_prev != batch_size or total_data_prev != total_data or (
batch_size >= total_data and sim_prev != sim):
print('Reading all data samples used in training...')
train_image, train_label, _, _, _ = get_data(dataset,
total_data,
dataset_file_path,
sim_round=sim)
batch_size_prev = batch_size
total_data_prev = total_data
sim_prev = sim
if batch_size >= total_data:
sampler = None
train_indices = indices_this_node
else:
sampler = MinibatchSampling(indices_this_node, batch_size, sim)
train_indices = None # To be defined later
last_batch_read_count = None
data_size_local = len(indices_this_node)