# UNPACK CONFIGS

(flag_para_load, train_filenames, val_filenames,

train_labels, val_labels, img_mean) = \

unpack_configs(config, ext_data=private_config['ext_data'],

ext_label=private_config['ext_label'])

gpu_send_queue = private_config['queue_gpu_send']

gpu_recv_queue = private_config['queue_gpu_recv']

# pycuda and zmq set up

drv.init()

dev = drv.Device(int(private_config['gpu'][-1]))

ctx = dev.make_context()

sock_gpu = zmq.Context().socket(zmq.PAIR)

if private_config['flag_client']:

sock_gpu.connect('tcp://localhost:{0}'.format(config['sock_gpu']))

else:

sock_gpu.bind('tcp://*:{0}'.format(config['sock_gpu']))

if flag_para_load:

sock_data = zmq.Context().socket(zmq.PAIR)

sock_data.connect('tcp://localhost:{0}'.format(

private_config['sock_data']))

load_send_queue = private_config['queue_t2l']

load_recv_queue = private_config['queue_l2t']

else:

load_send_queue = None

load_recv_queue = None

import theano.sandbox.cuda

theano.sandbox.cuda.use(private_config['gpu'])

import theano

theano.config.on_unused_input = 'warn'

from layers import DropoutLayer

from alex_net import AlexNet, compile_models

import theano.misc.pycuda_init

import theano.misc.pycuda_utils

## BUILD NETWORK ##

model = AlexNet(config)

layers = model.layers

batch_size = model.batch_size

## COMPILE FUNCTIONS ##

(train_model, validate_model, train_error, learning_rate,

shared_x, shared_y, rand_arr, vels) = compile_models(model, config)

total_params = model.params + vels

# total_params = model.params

# initialize gpuarrays that points to the theano shared variable

# pass parameters and other stuff

param_ga_list = []

param_other_list = []

param_ga_other_list = []

h_list = []

shape_list = []

dtype_list = []

average_fun_list = []

for param in total_params:

param_other = theano.shared(param.get_value())

param_ga = \

theano.misc.pycuda_utils.to_gpuarray(param.container.value)

param_ga_other = \

theano.misc.pycuda_utils.to_gpuarray(

param_other.container.value)

h = drv.mem_get_ipc_handle(param_ga.ptr)

average_fun = \

theano.function([], updates=[(param,

(param + param_other) / 2.)])

param_other_list.append(param_other)

param_ga_list.append(param_ga)

param_ga_other_list.append(param_ga_other)

h_list.append(h)

shape_list.append(param_ga.shape)

dtype_list.append(param_ga.dtype)

average_fun_list.append(average_fun)

# pass shape, dtype and handles

sock_gpu.send_pyobj((shape_list, dtype_list, h_list))

shape_other_list, dtype_other_list, h_other_list = sock_gpu.recv_pyobj()

param_ga_remote_list = []

# create gpuarray point to the other gpu use the passed information

for shape_other, dtype_other, h_other in zip(shape_other_list,

dtype_other_list,

h_other_list):

param_ga_remote = \

gpuarray.GPUArray(shape_other, dtype_other,

gpudata=drv.IPCMemoryHandle(h_other))

param_ga_remote_list.append(param_ga_remote)

print "Information passed between 2 GPUs"

##########################################

######################### TRAIN MODEL ################################

print '... training'

if flag_para_load:

# pass ipc handle and related information

gpuarray_batch = theano.misc.pycuda_utils.to_gpuarray(

shared_x.container.value)

h = drv.mem_get_ipc_handle(gpuarray_batch.ptr)

sock_data.send_pyobj((gpuarray_batch.shape, gpuarray_batch.dtype, h))

load_send_queue.put(img_mean)

n_train_batches = len(train_filenames)

minibatch_range = range(n_train_batches)

# gpu sync before start

gpu_send_queue.put('before_start')

assert gpu_recv_queue.get() == 'before_start'

# Start Training Loop

epoch = 0

step_idx = 0

val_record = []

while epoch < config['n_epochs']:

epoch = epoch + 1

if config['shuffle']:

np.random.shuffle(minibatch_range)

if config['resume_train'] and epoch == 1:

load_epoch = config['load_epoch']

load_weights(layers, config['weights_dir'], load_epoch)

epoch = load_epoch + 1

lr_to_load = np.load(

config['weights_dir'] + 'lr_' + str(load_epoch) + '.npy')

val_record = list(

np.load(config['weights_dir'] + 'val_record.npy'))

learning_rate.set_value(lr_to_load)

load_momentums(vels, config['weights_dir'], epoch)

if flag_para_load:

# send the initial message to load data, before each epoch

load_send_queue.put(str(train_filenames[minibatch_range[0]]))

load_send_queue.put(get_rand3d())

# clear the sync before 1st calc

load_send_queue.put('calc_finished')

count = 0

for minibatch_index in minibatch_range:

num_iter = (epoch - 1) * n_train_batches + count

count = count + 1

if count%20 == 1:

s = time.time()

cost_ij = train_model_wrap(train_model, shared_x,

shared_y, rand_arr, img_mean,

count, minibatch_index,

minibatch_range, batch_size,

train_filenames, train_labels,

flag_para_load,

config['batch_crop_mirror'],

send_queue=load_send_queue,

recv_queue=load_recv_queue)

# gpu sync

gpu_send_queue.put('after_train')

assert gpu_recv_queue.get() == 'after_train'

# exchanging weights

for param_ga, param_ga_other, param_ga_remote in \

zip(param_ga_list, param_ga_other_list,

param_ga_remote_list):

drv.memcpy_peer(param_ga_other.ptr,

param_ga_remote.ptr,

param_ga_remote.dtype.itemsize *

param_ga_remote.size,

ctx, ctx)

ctx.synchronize()

# gpu sync

gpu_send_queue.put('after_ctx_sync')

assert gpu_recv_queue.get() == 'after_ctx_sync'

# do average

for average_fun in average_fun_list:

average_fun()

# report train stats

if num_iter % config['print_freq'] == 0:

gpu_send_queue.put(cost_ij)

that_cost = gpu_recv_queue.get()

cost_ij = (cost_ij + that_cost) / 2.

if private_config['flag_verbose']:

print 'training @ iter = ', num_iter

print 'training cost:', cost_ij

if config['print_train_error']:

error_ij = train_error()

gpu_send_queue.put(error_ij)

that_error = gpu_recv_queue.get()

error_ij = (error_ij + that_error) / 2.

if private_config['flag_verbose']:

print 'training error rate:', error_ij

if flag_para_load and (count < len(minibatch_range)):

load_send_queue.put('calc_finished')

if count%20 == 0:

e = time.time()

print "time per 20 iter:", (e - s)

############### Test on Validation Set ##################

DropoutLayer.SetDropoutOff()

this_val_error, this_val_loss = get_val_error_loss(

rand_arr, shared_x, shared_y,

val_filenames, val_labels,

flag_para_load, img_mean,

batch_size, validate_model,

send_queue=load_send_queue, recv_queue=load_recv_queue)

# report validation stats

gpu_send_queue.put(this_val_error)

that_val_error = gpu_recv_queue.get()

this_val_error = (this_val_error + that_val_error) / 2.

gpu_send_queue.put(this_val_loss)

that_val_loss = gpu_recv_queue.get()

this_val_loss = (this_val_loss + that_val_loss) / 2.

if private_config['flag_verbose']:

print('epoch %i: validation loss %f ' %

(epoch, this_val_loss))

print('epoch %i: validation error %f %%' %

(epoch, this_val_error * 100.))

val_record.append([this_val_error, this_val_loss])

if private_config['flag_save']:

np.save(config['weights_dir'] + 'val_record.npy', val_record)

DropoutLayer.SetDropoutOn()

############################################

# Adapt Learning Rate

step_idx = adjust_learning_rate(config, epoch, step_idx,

val_record, learning_rate)

# Save Weights, only one of them will do

if private_config['flag_save']:

if epoch % config['snapshot_freq'] == 0:

save_weights(layers, config['weights_dir'], epoch)

np.save(config['weights_dir'] + 'lr_' + str(epoch) + '.npy',

learning_rate.get_value())

save_momentums(vels, config['weights_dir'], epoch)