def _factorized_reduction(self, x, out_filters, stride, is_training): """Reduces the shape of x without information loss due to striding.""" assert out_filters % 2 == 0, ( "Need even number of filters when using this factorized reduction.") if stride == 1: with tf.variable_scope("path_conv"): inp_c = self._get_C(x) w = create_weight("w", [1, 1, inp_c, out_filters]) x = tf.nn.conv2d(x, w, [1, 1, 1, 1], "SAME", data_format=self.data_format) x = batch_norm(x, is_training, data_format=self.data_format) return x stride_spec = self._get_strides(stride) # Skip path 1 path1 = tf.nn.avg_pool( x, [1, 1, 1, 1], stride_spec, "VALID", data_format=self.data_format) with tf.variable_scope("path1_conv"): inp_c = self._get_C(path1) w = create_weight("w", [1, 1, inp_c, out_filters // 2]) path1 = tf.nn.conv2d(path1, w, [1, 1, 1, 1], "SAME", data_format=self.data_format) # Skip path 2 # First pad with 0"s on the right and bottom, then shift the filter to # include those 0"s that were added. if self.data_format == "NHWC": pad_arr = [[0, 0], [0, 1], [0, 1], [0, 0]] path2 = tf.pad(x, pad_arr)[:, 1:, 1:, :] concat_axis = 3 else: pad_arr = [[0, 0], [0, 0], [0, 1], [0, 1]] path2 = tf.pad(x, pad_arr)[:, :, 1:, 1:] concat_axis = 1 path2 = tf.nn.avg_pool( path2, [1, 1, 1, 1], stride_spec, "VALID", data_format=self.data_format) with tf.variable_scope("path2_conv"): inp_c = self._get_C(path2) w = create_weight("w", [1, 1, inp_c, out_filters // 2]) path2 = tf.nn.conv2d(path2, w, [1, 1, 1, 1], "SAME", data_format=self.data_format) # Concat and apply BN final_path = tf.concat(values=[path1, path2], axis=concat_axis) final_path = batch_norm(final_path, is_training, data_format=self.data_format) return final_path
def _model(self, images, is_training, reuse=False): with tf.variable_scope(self.name, reuse=reuse): layers = [] out_filters = self.out_filters with tf.variable_scope("stem_conv"): w = create_weight("w", [3, 3, 3, out_filters]) x = tf.nn.conv2d( images, w, [1, 1, 1, 1], "SAME", data_format=self.data_format) x = batch_norm(x, is_training, data_format=self.data_format) layers.append(x) if self.whole_channels: start_idx = 0 else: start_idx = self.num_branches for layer_id in range(self.num_layers): with tf.variable_scope("layer_{0}".format(layer_id)): if self.fixed_arc is None: x = self._enas_layer( layer_id, layers, start_idx, out_filters, is_training) else: x = self._fixed_layer( layer_id, layers, start_idx, out_filters, is_training) layers.append(x) if layer_id in self.pool_layers: if self.fixed_arc is not None: out_filters *= 2 with tf.variable_scope("pool_at_{0}".format(layer_id)): pooled_layers = [] for i, layer in enumerate(layers): with tf.variable_scope("from_{0}".format(i)): x = self._factorized_reduction( layer, out_filters, 2, is_training) pooled_layers.append(x) layers = pooled_layers if self.whole_channels: start_idx += 1 + layer_id else: start_idx += 2 * self.num_branches + layer_id print(layers[-1]) x = global_avg_pool(x, data_format=self.data_format) if is_training: x = tf.nn.dropout(x, self.keep_prob) with tf.variable_scope("fc"): if self.data_format == "NWHC": inp_c = x.get_shape()[3].value elif self.data_format == "NCHW": inp_c = x.get_shape()[1].value else: raise ValueError( "Unknown data_format {0}".format(self.data_format)) w = create_weight("w", [inp_c, 10]) x = tf.matmul(x, w) return x
def _maybe_calibrate_size(self, layers, out_filters, is_training): """Makes sure layers[0] and layers[1] have the same shapes.""" hw = [self._get_HW(layer) for layer in layers] c = [self._get_C(layer) for layer in layers] with tf.variable_scope("calibrate"): x = layers[0] if hw[0] != hw[1]: assert hw[0] == 2 * hw[1] with tf.variable_scope("pool_x"): x = tf.nn.relu(x) x = self._factorized_reduction(x, out_filters, 2, is_training) elif c[0] != out_filters: with tf.variable_scope("pool_x"): w = create_weight("w", [1, 1, c[0], out_filters]) x = tf.nn.relu(x) x = tf.nn.conv2d(x, w, [1, 1, 1, 1], "SAME", data_format=self.data_format) x = batch_norm(x, is_training, data_format=self.data_format) y = layers[1] if c[1] != out_filters: with tf.variable_scope("pool_y"): w = create_weight("w", [1, 1, c[1], out_filters]) y = tf.nn.relu(y) y = tf.nn.conv2d(y, w, [1, 1, 1, 1], "SAME", data_format=self.data_format) y = batch_norm(y, is_training, data_format=self.data_format) return [x, y]
def _pool_branch(self, inputs, is_training, count, avg_or_max, start_idx=None): """ Args: start_idx: where to start taking the output channels. if None, assuming fixed_arc mode count: how many output_channels to take. """ if start_idx is None: assert self.fixed_arc is not None, "you screwed up!" if self.data_format == "NHWC": inp_c = inputs.get_shape()[3].value elif self.data_format == "NCHW": inp_c = inputs.get_shape()[1].value with tf.variable_scope("conv_1"): w = create_weight("w", [1, 1, inp_c, self.out_filters]) x = tf.nn.conv2d(inputs, w, [1, 1, 1, 1], "SAME", data_format=self.data_format) x = batch_norm(x, is_training, data_format=self.data_format) x = tf.nn.relu(x) with tf.variable_scope("pool"): if self.data_format == "NHWC": actual_data_format = "channels_last" elif self.data_format == "NCHW": actual_data_format = "channels_first" if avg_or_max == "avg": x = tf.layers.average_pooling2d( x, [3, 3], [1, 1], "SAME", data_format=actual_data_format) elif avg_or_max == "max": x = tf.layers.max_pooling2d( x, [3, 3], [1, 1], "SAME", data_format=actual_data_format) else: raise ValueError("Unknown pool {}".format(avg_or_max)) if start_idx is not None: if self.data_format == "NHWC": x = x[:, :, :, start_idx: start_idx+count] elif self.data_format == "NCHW": x = x[:, start_idx: start_idx+count, :, :] return x
def _fixed_conv(self, x, f_size, out_filters, stride, is_training, stack_convs=2): """ Apply fixed separable convolution. Args: stacked_convs: number of separable convs to apply. """ for conv_id in range(stack_convs): inp_c = self._get_C(x) if conv_id == 0: strides = self._get_strides(stride) else: strides = [1, 1, 1, 1] with tf.variable_scope("sep_conv_{}".format(conv_id)): w_depthwise = create_weight("w_depth", [f_size, f_size, inp_c, 1]) w_pointwise = create_weight("w_point", [1, 1, inp_c, out_filters]) x = tf.nn.relu(x) x = tf.nn.separable_conv2d(x, depthwise_filter=w_depthwise, pointwise_filter=w_pointwise, strides=strides, padding="SAME", data_format=self.data_format) x = batch_norm(x, is_training, data_format=self.data_format) return x
def _conv_branch(self, inputs, filter_size, is_training, count, out_filters, ch_mul=1, start_idx=None, separable=False): """ Args: start_idx: where to start taking the output channels. if None, assuming fixed_arc mode count: how many output_channels to take. """ if start_idx is None: assert self.fixed_arc is not None, "you screwed up!" if self.data_format == "NHWC": inp_c = inputs.get_shape()[3].value elif self.data_format == "NCHW": inp_c = inputs.get_shape()[1].value with tf.variable_scope("inp_conv_1"): w = create_weight("w", [1, 1, inp_c, out_filters]) x = tf.nn.conv2d(inputs, w, [1, 1, 1, 1], "SAME", data_format=self.data_format) x = batch_norm(x, is_training, data_format=self.data_format) x = tf.nn.relu(x) with tf.variable_scope("out_conv_{}".format(filter_size)): if start_idx is None: if separable: w_depth = create_weight( "w_depth", [self.filter_size, self.filter_size, out_filters, ch_mul]) w_point = create_weight( "w_point", [1, 1, out_filters * ch_mul, count]) x = tf.nn.separable_conv2d(x, w_depth, w_point, strides=[1, 1, 1, 1], padding="SAME", data_format=self.data_format) x = batch_norm( x, is_training, data_format=self.data_format) else: w = create_weight( "w", [filter_size, filter_size, inp_c, count]) x = tf.nn.conv2d( x, w, [1, 1, 1, 1], "SAME", data_format=self.data_format) x = batch_norm( x, is_training, data_format=self.data_format) else: if separable: w_depth = create_weight( "w_depth", [filter_size, filter_size, out_filters, ch_mul]) w_point = create_weight( "w_point", [out_filters, out_filters * ch_mul]) w_point = w_point[start_idx:start_idx+count, :] w_point = tf.transpose(w_point, [1, 0]) w_point = tf.reshape( w_point, [1, 1, out_filters * ch_mul, count]) x = tf.nn.separable_conv2d(x, w_depth, w_point, strides=[1, 1, 1, 1], padding="SAME", data_format=self.data_format) mask = tf.range(0, out_filters, dtype=tf.int32) mask = tf.logical_and( start_idx <= mask, mask < start_idx + count) x = batch_norm_with_mask( x, is_training, mask, out_filters, data_format=self.data_format) else: w = create_weight( "w", [filter_size, filter_size, out_filters, out_filters]) w = tf.transpose(w, [3, 0, 1, 2]) w = w[start_idx:start_idx+count, :, :, :] w = tf.transpose(w, [1, 2, 3, 0]) x = tf.nn.conv2d( x, w, [1, 1, 1, 1], "SAME", data_format=self.data_format) mask = tf.range(0, out_filters, dtype=tf.int32) mask = tf.logical_and( start_idx <= mask, mask < start_idx + count) x = batch_norm_with_mask( x, is_training, mask, out_filters, data_format=self.data_format) x = tf.nn.relu(x) return x
def _fixed_layer( self, layer_id, prev_layers, start_idx, out_filters, is_training): """ Args: layer_id: current layer prev_layers: cache of previous layers. for skip connections start_idx: where to start looking at. technically, we can infer this from layer_id, but why bother... is_training: for batch_norm """ inputs = prev_layers[-1] if self.whole_channels: if self.data_format == "NHWC": inp_c = inputs.get_shape()[3].value elif self.data_format == "NCHW": inp_c = inputs.get_shape()[1].value count = self.sample_arc[start_idx] if count in [0, 1, 2, 3]: size = [3, 3, 5, 5] filter_size = size[count] with tf.variable_scope("conv_1x1"): w = create_weight("w", [1, 1, inp_c, out_filters]) out = tf.nn.relu(inputs) out = tf.nn.conv2d(out, w, [1, 1, 1, 1], "SAME", data_format=self.data_format) out = batch_norm(out, is_training, data_format=self.data_format) with tf.variable_scope("conv_{0}x{0}".format(filter_size)): w = create_weight( "w", [filter_size, filter_size, out_filters, out_filters]) out = tf.nn.relu(out) out = tf.nn.conv2d(out, w, [1, 1, 1, 1], "SAME", data_format=self.data_format) out = batch_norm(out, is_training, data_format=self.data_format) elif count == 4: pass elif count == 5: pass else: raise ValueError( "Unknown operation number '{0}'".format(count)) else: count = (self.sample_arc[start_idx:start_idx + 2*self.num_branches] * self.out_filters_scale) branches = [] total_out_channels = 0 with tf.variable_scope("branch_0"): total_out_channels += count[1] branches.append(self._conv_branch( inputs, 3, is_training, count[1])) with tf.variable_scope("branch_1"): total_out_channels += count[3] branches.append( self._conv_branch(inputs, 3, is_training, count[3], separable=True)) with tf.variable_scope("branch_2"): total_out_channels += count[5] branches.append(self._conv_branch( inputs, 5, is_training, count[5])) with tf.variable_scope("branch_3"): total_out_channels += count[7] branches.append( self._conv_branch(inputs, 5, is_training, count[7], separable=True)) if self.num_branches >= 5: with tf.variable_scope("branch_4"): total_out_channels += count[9] branches.append( self._pool_branch(inputs, is_training, count[9], "avg")) if self.num_branches >= 6: with tf.variable_scope("branch_5"): total_out_channels += count[11] branches.append( self._pool_branch(inputs, is_training, count[11], "max")) with tf.variable_scope("final_conv"): w = create_weight("w", [1, 1, total_out_channels, out_filters]) if self.data_format == "NHWC": branches = tf.concat(branches, axis=3) elif self.data_format == "NCHW": branches = tf.concat(branches, axis=1) out = tf.nn.relu(branches) out = tf.nn.conv2d(out, w, [1, 1, 1, 1], "SAME", data_format=self.data_format) out = batch_norm(out, is_training, data_format=self.data_format) if layer_id > 0: if self.whole_channels: skip_start = start_idx + 1 else: skip_start = start_idx + 2 * self.num_branches skip = self.sample_arc[skip_start: skip_start + layer_id] total_skip_channels = np.sum(skip) + 1 res_layers = [] for i in range(layer_id): if skip[i] == 1: res_layers.append(prev_layers[i]) prev = res_layers + [out] if self.data_format == "NHWC": prev = tf.concat(prev, axis=3) elif self.data_format == "NCHW": prev = tf.concat(prev, axis=1) out = prev with tf.variable_scope("skip"): w = create_weight( "w", [1, 1, total_skip_channels * out_filters, out_filters]) out = tf.nn.relu(out) out = tf.nn.conv2d( out, w, [1, 1, 1, 1], "SAME", data_format=self.data_format) out = batch_norm(out, is_training, data_format=self.data_format) return out
def _enas_layer(self, layer_id, prev_layers, start_idx, out_filters, is_training): """ Args: layer_id: current layer prev_layers: cache of previous layers. for skip connections start_idx: where to start looking at. technically, we can infer this from layer_id, but why bother... is_training: for batch_norm """ inputs = prev_layers[-1] if self.whole_channels: if self.data_format == "NHWC": inp_h = inputs.get_shape()[1].value inp_w = inputs.get_shape()[2].value inp_c = inputs.get_shape()[3].value elif self.data_format == "NCHW": inp_c = inputs.get_shape()[1].value inp_h = inputs.get_shape()[2].value inp_w = inputs.get_shape()[3].value count = self.sample_arc[start_idx] branches = {} with tf.variable_scope("branch_0"): y = self._conv_branch(inputs, 3, is_training, out_filters, out_filters, start_idx=0) branches[tf.equal(count, 0)] = lambda: y with tf.variable_scope("branch_1"): y = self._conv_branch(inputs, 3, is_training, out_filters, out_filters, start_idx=0, separable=True) branches[tf.equal(count, 1)] = lambda: y with tf.variable_scope("branch_2"): y = self._conv_branch(inputs, 5, is_training, out_filters, out_filters, start_idx=0) branches[tf.equal(count, 2)] = lambda: y with tf.variable_scope("branch_3"): y = self._conv_branch(inputs, 5, is_training, out_filters, out_filters, start_idx=0, separable=True) branches[tf.equal(count, 3)] = lambda: y if self.num_branches >= 5: with tf.variable_scope("branch_4"): y = self._pool_branch(inputs, is_training, out_filters, "avg", start_idx=0) branches[tf.equal(count, 4)] = lambda: y if self.num_branches >= 6: with tf.variable_scope("branch_5"): y = self._pool_branch(inputs, is_training, out_filters, "max", start_idx=0) branches[tf.equal(count, 5)] = lambda: y out = tf.case(branches, default=lambda: tf.constant(0, tf.float32), exclusive=True) if self.data_format == "NHWC": out.set_shape([None, inp_h, inp_w, out_filters]) elif self.data_format == "NCHW": out.set_shape([None, out_filters, inp_h, inp_w]) else: count = self.sample_arc[start_idx:start_idx + 2 * self.num_branches] branches = [] with tf.variable_scope("branch_0"): branches.append(self._conv_branch(inputs, 3, is_training, count[1], out_filters, start_idx=count[0])) with tf.variable_scope("branch_1"): branches.append(self._conv_branch(inputs, 3, is_training, count[3], out_filters, start_idx=count[2], separable=True)) with tf.variable_scope("branch_2"): branches.append(self._conv_branch(inputs, 5, is_training, count[5], out_filters, start_idx=count[4])) with tf.variable_scope("branch_3"): branches.append(self._conv_branch(inputs, 5, is_training, count[7], out_filters, start_idx=count[6], separable=True)) if self.num_branches >= 5: with tf.variable_scope("branch_4"): branches.append(self._pool_branch(inputs, is_training, count[9], "avg", start_idx=count[8])) if self.num_branches >= 6: with tf.variable_scope("branch_5"): branches.append(self._pool_branch(inputs, is_training, count[11], "max", start_idx=count[10])) with tf.variable_scope("final_conv"): w = create_weight( "w", [self.num_branches * out_filters, out_filters]) w_mask = tf.constant( [False] * (self.num_branches * out_filters), tf.bool) new_range = tf.range( 0, self.num_branches * out_filters, dtype=tf.int32) for i in range(self.num_branches): start = out_filters * i + count[2 * i] new_mask = tf.logical_and( start <= new_range, new_range < start + count[2 * i + 1]) w_mask = tf.logical_or(w_mask, new_mask) w = tf.boolean_mask(w, w_mask) w = tf.reshape(w, [1, 1, -1, out_filters]) inp = prev_layers[-1] if self.data_format == "NHWC": branches = tf.concat(branches, axis=3) elif self.data_format == "NCHW": branches = tf.concat(branches, axis=1) N = tf.shape(inp)[0] H = inp.get_shape()[2].value W = inp.get_shape()[3].value branches = tf.reshape(branches, [N, -1, H, W]) out = tf.nn.conv2d( branches, w, [1, 1, 1, 1], "SAME", data_format=self.data_format) out = batch_norm(out, is_training, data_format=self.data_format) out = tf.nn.relu(out) if layer_id > 0: if self.whole_channels: skip_start = start_idx + 1 else: skip_start = start_idx + 2 * self.num_branches skip = self.sample_arc[skip_start: skip_start + layer_id] with tf.variable_scope("skip"): res_layers = [] for i in range(layer_id): res_layers.append(tf.cond(tf.equal(skip[i], 1), lambda: prev_layers[i], lambda: tf.zeros_like(prev_layers[i]))) res_layers.append(out) out = tf.add_n(res_layers) out = batch_norm(out, is_training, data_format=self.data_format) return out
def _enas_layer(self, layer_id, prev_layers, arc, out_filters): """ Args: layer_id: current layer prev_layers: cache of previous layers. for skip connections start_idx: where to start looking at. technically, we can infer this from layer_id, but why bother... """ assert len(prev_layers) == 2, "need exactly 2 inputs" layers = [prev_layers[0], prev_layers[1]] layers = self._maybe_calibrate_size( layers, out_filters, is_training=True ) # make the shape layer[0] same to shape of layer[1] used = [] for cell_id in range(self.num_cells): prev_layers = tf.stack(layers, axis=0) with tf.variable_scope("cell_{0}".format(cell_id)): with tf.variable_scope("x"): x_id = arc[4 * cell_id] x_op = arc[4 * cell_id + 1] x = prev_layers[x_id, :, :, :, :] x = self._enas_cell( x, cell_id, x_id, x_op, out_filters ) # Get the result of the op_id corresponding to the specified x_id x_used = tf.one_hot(x_id, depth=self.num_cells + 2, dtype=tf.int32) with tf.variable_scope("y"): y_id = arc[4 * cell_id + 2] y_op = arc[4 * cell_id + 3] y = prev_layers[y_id, :, :, :, :] y = self._enas_cell(y, cell_id, y_id, y_op, out_filters) y_used = tf.one_hot(y_id, depth=self.num_cells + 2, dtype=tf.int32) out = x + y used.extend([x_used, y_used]) layers.append(out) used = tf.add_n(used) # Wise addition of all elements in "used" # Find the coordinates of the element equal to 0 in used, and return a few coordinates when there are several, the dimension of each coordinate is the same as the dimension of used Question: why do this? indices = tf.where(tf.equal(used, 0)) indices = tf.to_int32(indices) indices = tf.reshape(indices, [-1]) num_outs = tf.size(indices) out = tf.stack(layers, axis=0) out = tf.gather(out, indices, axis=0) # question: what's for? inp = prev_layers[0] if self.data_format == "NHWC": N = tf.shape(inp)[0] H = tf.shape(inp)[1] W = tf.shape(inp)[2] C = tf.shape(inp)[3] out = tf.transpose(out, [1, 2, 3, 0, 4]) out = tf.reshape(out, [N, H, W, num_outs * out_filters]) elif self.data_format == "NCHW": N = tf.shape(inp)[0] C = tf.shape(inp)[1] H = tf.shape(inp)[2] W = tf.shape(inp)[3] out = tf.transpose(out, [1, 0, 2, 3, 4]) out = tf.reshape(out, [N, num_outs * out_filters, H, W]) else: raise ValueError("Unknown data_format '{0}'".format( self.data_format)) with tf.variable_scope("final_conv"): w = create_weight("w", [self.num_cells + 2, out_filters * out_filters]) w = tf.gather(w, indices, axis=0) w = tf.reshape(w, [1, 1, num_outs * out_filters, out_filters]) out = tf.nn.relu(out) out = tf.nn.conv2d(out, w, strides=[1, 1, 1, 1], padding="SAME", data_format=self.data_format) out = batch_norm(out, is_training=True, data_format=self.data_format) out = tf.reshape(out, tf.shape(prev_layers[0])) return out
def _enas_cell(self, x, curr_cell, prev_cell, op_id, out_filters): """Performs an enas operation specified by op_id.""" num_possible_inputs = curr_cell + 1 with tf.variable_scope("avg_pool"): avg_pool = tf.layers.average_pooling2d( x, [3, 3], [1, 1], "SAME", data_format=self.actual_data_format) avg_pool_c = self._get_C(avg_pool) if avg_pool_c != out_filters: with tf.variable_scope("conv"): w = create_weight( "w", [num_possible_inputs, avg_pool_c * out_filters]) w = w[prev_cell] w = tf.reshape(w, [1, 1, avg_pool_c, out_filters]) avg_pool = tf.nn.relu(avg_pool) avg_pool = tf.nn.conv2d(avg_pool, w, strides=[1, 1, 1, 1], padding="SAME", data_format=self.data_format) avg_pool = batch_norm(avg_pool, is_training=True, data_format=self.data_format) with tf.variable_scope("max_pool"): max_pool = tf.layers.max_pooling2d( x, [3, 3], [1, 1], "SAME", data_format=self.actual_data_format) max_pool_c = self._get_C(max_pool) if max_pool_c != out_filters: with tf.variable_scope("conv"): w = create_weight( "w", [num_possible_inputs, max_pool_c * out_filters]) w = w[prev_cell] w = tf.reshape(w, [1, 1, max_pool_c, out_filters]) max_pool = tf.nn.relu(max_pool) max_pool = tf.nn.conv2d(max_pool, w, strides=[1, 1, 1, 1], padding="SAME", data_format=self.data_format) max_pool = batch_norm(max_pool, is_training=True, data_format=self.data_format) # 获取通道数 x_c = self._get_C(x) if x_c != out_filters: with tf.variable_scope("x_conv"): w = create_weight("w", [num_possible_inputs, x_c * out_filters]) w = w[prev_cell] w = tf.reshape(w, [1, 1, x_c, out_filters]) x = tf.nn.relu(x) x = tf.nn.conv2d(x, w, strides=[1, 1, 1, 1], padding="SAME", data_format=self.data_format) x = batch_norm(x, is_training=True, data_format=self.data_format) out = [ self._enas_conv(x, curr_cell, prev_cell, 3, out_filters), # separable_conv2d 3*3 self._enas_conv(x, curr_cell, prev_cell, 5, out_filters), # separable_conv2d 5*5 avg_pool, max_pool, x, ] out = tf.stack(out, axis=0) out = out[op_id, :, :, :, :] # choose operation by op_id return out
def _fixed_layer(self, layer_id, prev_layers, arc, out_filters, stride, is_training, normal_or_reduction_cell="normal"): """ Args: prev_layers: cache of previous layers. for skip connections is_training: for batch_norm """ assert len(prev_layers) == 2 layers = [prev_layers[0], prev_layers[1]] layers = self._maybe_calibrate_size(layers, out_filters, is_training=is_training) with tf.variable_scope("layer_base"): x = layers[1] inp_c = self._get_C(x) w = create_weight("w", [1, 1, inp_c, out_filters]) x = tf.nn.relu(x) x = tf.nn.conv2d(x, w, [1, 1, 1, 1], "SAME", data_format=self.data_format) x = batch_norm(x, is_training, data_format=self.data_format) layers[1] = x used = np.zeros([self.num_cells + 2], dtype=np.int32) f_sizes = [3, 5] for cell_id in range(self.num_cells): with tf.variable_scope("cell_{}".format(cell_id)): x_id = arc[4 * cell_id] used[x_id] += 1 x_op = arc[4 * cell_id + 1] x = layers[x_id] x_stride = stride if x_id in [0, 1] else 1 with tf.variable_scope("x_conv"): if x_op in [0, 1]: f_size = f_sizes[x_op] x = self._fixed_conv(x, f_size, out_filters, x_stride, is_training) elif x_op in [2, 3]: inp_c = self._get_C(x) if x_op == 2: x = tf.layers.average_pooling2d( x, [3, 3], [x_stride, x_stride], "SAME", data_format=self.actual_data_format) else: x = tf.layers.max_pooling2d( x, [3, 3], [x_stride, x_stride], "SAME", data_format=self.actual_data_format) if inp_c != out_filters: w = create_weight("w", [1, 1, inp_c, out_filters]) x = tf.nn.relu(x) x = tf.nn.conv2d(x, w, [1, 1, 1, 1], "SAME", data_format=self.data_format) x = batch_norm(x, is_training, data_format=self.data_format) else: inp_c = self._get_C(x) if x_stride > 1: assert x_stride == 2 x = self._factorized_reduction( x, out_filters, 2, is_training) if inp_c != out_filters: w = create_weight("w", [1, 1, inp_c, out_filters]) x = tf.nn.relu(x) x = tf.nn.conv2d(x, w, [1, 1, 1, 1], "SAME", data_format=self.data_format) x = batch_norm(x, is_training, data_format=self.data_format) if (x_op in [0, 1, 2, 3] and self.drop_path_keep_prob is not None and is_training): x = self._apply_drop_path(x, layer_id) y_id = arc[4 * cell_id + 2] used[y_id] += 1 y_op = arc[4 * cell_id + 3] y = layers[y_id] y_stride = stride if y_id in [0, 1] else 1 with tf.variable_scope("y_conv"): if y_op in [0, 1]: f_size = f_sizes[y_op] y = self._fixed_conv(y, f_size, out_filters, y_stride, is_training) elif y_op in [2, 3]: inp_c = self._get_C(y) if y_op == 2: y = tf.layers.average_pooling2d( y, [3, 3], [y_stride, y_stride], "SAME", data_format=self.actual_data_format) else: y = tf.layers.max_pooling2d( y, [3, 3], [y_stride, y_stride], "SAME", data_format=self.actual_data_format) if inp_c != out_filters: w = create_weight("w", [1, 1, inp_c, out_filters]) y = tf.nn.relu(y) y = tf.nn.conv2d(y, w, [1, 1, 1, 1], "SAME", data_format=self.data_format) y = batch_norm(y, is_training, data_format=self.data_format) else: inp_c = self._get_C(y) if y_stride > 1: assert y_stride == 2 y = self._factorized_reduction( y, out_filters, 2, is_training) if inp_c != out_filters: w = create_weight("w", [1, 1, inp_c, out_filters]) y = tf.nn.relu(y) y = tf.nn.conv2d(y, w, [1, 1, 1, 1], "SAME", data_format=self.data_format) y = batch_norm(y, is_training, data_format=self.data_format) if (y_op in [0, 1, 2, 3] and self.drop_path_keep_prob is not None and is_training): y = self._apply_drop_path(y, layer_id) out = x + y layers.append(out) out = self._fixed_combine(layers, used, out_filters, is_training, normal_or_reduction_cell) return out
def _model(self, images, is_training, reuse=False): """Compute the logits given the images.""" if self.fixed_arc is None: is_training = True # is_training = True with tf.variable_scope(self.name, reuse=reuse): # the first two inputs with tf.variable_scope("stem_conv"): w = create_weight("w", [3, 3, 3, self.out_filters * 3]) x = tf.nn.conv2d(images, w, [1, 1, 1, 1], "SAME", data_format=self.data_format) x = batch_norm(x, is_training, data_format=self.data_format) if self.data_format == "NHCW": split_axis = 3 elif self.data_format == "NCHW": split_axis = 1 else: raise ValueError("Unknown data_format '{0}'".format( self.data_format)) layers = [x, x] # building layers in the micro space out_filters = self.out_filters for layer_id in range(self.num_layers + 2): with tf.variable_scope("layer_{0}".format(layer_id)): # generate normal cell if layer_id not in self.pool_layers: if self.fixed_arc is None: x = self._enas_layer(layer_id, layers, self.normal_arc, out_filters) else: x = self._fixed_layer( layer_id, layers, self.normal_arc, out_filters, 1, is_training, normal_or_reduction_cell="normal") # generate reduction cell else: out_filters *= 2 if self.fixed_arc is None: x = self._factorized_reduction( x, out_filters, 2, is_training) layers = [layers[-1], x] x = self._enas_layer(layer_id, layers, self.reduce_arc, out_filters) else: x = self._fixed_layer( layer_id, layers, self.reduce_arc, out_filters, 2, is_training, normal_or_reduction_cell="reduction") print("Layer {0:>2d}: {1}".format(layer_id, x)) layers = [layers[-1], x] # auxiliary heads ''' This is an auxiliary classification network. It can output the predicted classification results, and then the resulting loss will be collected into the total loss, thus improving the accuracy. ''' self.num_aux_vars = 0 if (self.use_aux_heads and layer_id in self.aux_head_indices and is_training): print("Using aux_head at layer {0}".format(layer_id)) with tf.variable_scope("aux_head"): aux_logits = tf.nn.relu(x) aux_logits = tf.layers.average_pooling2d( aux_logits, [5, 5], [3, 3], "VALID", data_format=self.actual_data_format) with tf.variable_scope("proj"): # question what's for? inp_c = self._get_C(aux_logits) w = create_weight("w", [1, 1, inp_c, 128]) aux_logits = tf.nn.conv2d( aux_logits, w, [1, 1, 1, 1], "SAME", data_format=self.data_format) aux_logits = batch_norm( aux_logits, is_training=True, data_format=self.data_format) aux_logits = tf.nn.relu(aux_logits) with tf.variable_scope("avg_pool"): inp_c = self._get_C(aux_logits) hw = self._get_HW(aux_logits) print(hw) w = create_weight("w", [hw, hw, inp_c, 768]) aux_logits = tf.nn.conv2d( aux_logits, w, [1, 1, 1, 1], "SAME", data_format=self.data_format) aux_logits = batch_norm( aux_logits, is_training=True, data_format=self.data_format) aux_logits = tf.nn.relu(aux_logits) with tf.variable_scope("fc"): aux_logits = global_avg_pool( aux_logits, data_format=self.data_format) inp_c = aux_logits.get_shape()[1].value w = create_weight("w", [inp_c, self.output_classes]) aux_logits = tf.matmul(aux_logits, w) self.aux_logits = aux_logits aux_head_variables = [ var for var in tf.trainable_variables() if (var.name.startswith(self.name) and "aux_head" in var.name) ] self.num_aux_vars = count_model_params(aux_head_variables) print("Aux head uses {0} params".format(self.num_aux_vars)) x = tf.nn.relu(x) x = global_avg_pool(x, data_format=self.data_format) if is_training and self.keep_prob is not None and self.keep_prob < 1.0: x = tf.nn.dropout(x, self.keep_prob) # FC layer classification with tf.variable_scope("fc"): inp_c = self._get_C(x) w = create_weight("w", [inp_c, self.output_classes]) x = tf.matmul(x, w) return x