diff --git a/video_prediction_savp/video_prediction/layers/layer_def.py b/video_prediction_savp/video_prediction/layers/layer_def.py index 6b7f4387001c9318507ad809d7176071312742d0..a59643c7a6d69141134ec01c9b147c4798bfed8e 100644 --- a/video_prediction_savp/video_prediction/layers/layer_def.py +++ b/video_prediction_savp/video_prediction/layers/layer_def.py @@ -18,7 +18,7 @@ def _activation_summary(x): tf.summary.histogram(tensor_name + '/activations', x) tf.summary.scalar(tensor_name + '/sparsity', tf.nn.zero_fraction(x)) -def _variable_on_cpu(name, shape, initializer): +def _variable_on_gpu(name, shape, initializer): """Helper to create a Variable stored on CPU memory. Args: name: name of the variable @@ -27,7 +27,7 @@ def _variable_on_cpu(name, shape, initializer): Returns: Variable Tensor """ - with tf.device('/cpu:0'): + with tf.device('/gpu:0'): var = tf.get_variable(name, shape, initializer=initializer) return var @@ -45,8 +45,8 @@ def _variable_with_weight_decay(name, shape, stddev, wd,initializer=tf.contrib.l Returns: Variable Tensor """ - #var = _variable_on_cpu(name, shape,tf.truncated_normal_initializer(stddev = stddev)) - var = _variable_on_cpu(name, shape, initializer) + #var = _variable_on_gpu(name, shape,tf.truncated_normal_initializer(stddev = stddev)) + var = _variable_on_gpu(name, shape, initializer) if wd: weight_decay = tf.multiply(tf.nn.l2_loss(var), wd, name = 'weight_loss') weight_decay.set_shape([]) @@ -63,7 +63,7 @@ def conv_layer(inputs, kernel_size, stride, num_features, idx, initializer=tf.co weights = _variable_with_weight_decay('weights',shape = [kernel_size, kernel_size, input_channels, num_features], stddev = 0.01, wd = weight_decay) - biases = _variable_on_cpu('biases', [num_features], initializer) + biases = _variable_on_gpu('biases', [num_features], initializer) conv = tf.nn.conv2d(inputs, weights, strides = [1, stride, stride, 1], padding = 'SAME') conv_biased = tf.nn.bias_add(conv, biases) if activate == "linear": @@ -88,7 +88,7 @@ def transpose_conv_layer(inputs, kernel_size, stride, num_features, idx, initial weights = _variable_with_weight_decay('weights', shape = [kernel_size, kernel_size, num_features, input_channels], stddev = 0.1, wd = weight_decay) - biases = _variable_on_cpu('biases', [num_features],initializer) + biases = _variable_on_gpu('biases', [num_features],initializer) batch_size = tf.shape(inputs)[0] output_shape = tf.stack( @@ -122,7 +122,7 @@ def fc_layer(inputs, hiddens, idx, flat=False, activate="relu",weight_init=0.01, weights = _variable_with_weight_decay('weights', shape = [dim, hiddens], stddev = weight_init, wd = weight_decay) - biases = _variable_on_cpu('biases', [hiddens],initializer) + biases = _variable_on_gpu('biases', [hiddens],initializer) if activate == "linear": return tf.add(tf.matmul(inputs_processed, weights), biases, name = str(idx) + '_fc') elif activate == "sigmoid":