Actualizo master con el long() y separo en dos simplify_network_def_based_on_constraint

.gitignore

@@ -2,4 +2,4 @@ data/*
 models/*
 **/__pycache__/
 outs/
-
+**.out

functions.py

@@ -524,6 +524,58 @@ def build_latency_lookup_table(network_def_full, lookup_table_path, min_conv_fea
         pickle.dump(lookup_table, file_id)      
     return 
 
+def compose_simplified_network_def(simplified_network_def, network_def, current_num_out_channels, target_layer_indices, max_num_out_channels):
+    for target_layer_index in target_layer_indices:
+        update_num_out_channels = True
+        current_num_out_channels_after_pixel_shuffle = current_num_out_channels
+        for layer_idx, (layer_name, layer_properties) in enumerate(simplified_network_def.items()):
+            if layer_idx < target_layer_index:
+                continue
+            
+            # for the block to be simplified (# of output channels is simplified)
+            if update_num_out_channels:
+                if not layer_properties[KEY_IS_DEPTHWISE]:
+                    layer_properties[KEY_NUM_OUT_CHANNELS] = current_num_out_channels
+                    update_num_out_channels = False
+                    
+                    print('    simplify_def>     layer {}: num of output channel changed to {}'.format(layer_name, str(current_num_out_channels)))
+                else:
+                    raise ValueError('Expected a non-depthwise layer but got a depthwise layer.')
+            # for blocks following the target blocks (# of input channels is simplified)
+            else:
+                if current_num_out_channels_after_pixel_shuffle % layer_properties[
+                    KEY_BEFORE_SQUARED_PIXEL_SHUFFLE_FACTOR] != 0:
+                    raise ValueError('current_num_out_channels or current_num_out_channels_after_pixel_shuffle is '
+                                        'not divisible by the scaling factor of pixel shuffling.')
+                current_num_out_channels_after_pixel_shuffle = (
+                        current_num_out_channels_after_pixel_shuffle / layer_properties[
+                    KEY_BEFORE_SQUARED_PIXEL_SHUFFLE_FACTOR])
+                layer_properties[KEY_NUM_IN_CHANNELS] = current_num_out_channels_after_pixel_shuffle
+                print('    simplify_def>     layer {}: num of input channel changed to {}'.format(layer_name, str(current_num_out_channels_after_pixel_shuffle)))
+
+                '''
+                    Consider the case that a FC layer is placed after a Conv and Flatten:
+                        FC: input feature size: Cin
+                            output feature size: Cout
+                        Conv: output feature map size: H x W x C
+                        So Cin = H x W x C.
+                        If C -> C' based on constraints, then Cin -> H x W x C'
+                '''
+                if layer_properties[KEY_BEFORE_SQUARED_PIXEL_SHUFFLE_FACTOR] == 1:
+                    if network_def[layer_name][KEY_NUM_IN_CHANNELS] > max_num_out_channels:
+                        assert network_def[layer_name][KEY_NUM_IN_CHANNELS] % max_num_out_channels == 0
+                        # H x W here
+                        spatial_factor = network_def[layer_name][KEY_NUM_IN_CHANNELS] // max_num_out_channels
+                        layer_properties[KEY_NUM_IN_CHANNELS] = spatial_factor*current_num_out_channels
+                        print('    simplify_def>     [Update] layer {}: num of input channel changed to {}'.format(layer_name, str(spatial_factor*current_num_out_channels)))
+
+                if not layer_properties[KEY_IS_DEPTHWISE]:
+                    break
+                else:
+                    layer_properties[KEY_NUM_OUT_CHANNELS] = current_num_out_channels_after_pixel_shuffle
+                    layer_properties[KEY_GROUPS] = current_num_out_channels_after_pixel_shuffle
+                    print('    simplify_def>     depthwise layer {}: num of output channel changed to {}'.format(layer_name, str(current_num_out_channels_after_pixel_shuffle)))
+    return simplified_network_def
 
 def simplify_network_def_based_on_constraint(network_def, block, constraint, resource_type,
                                              lookup_table_path=None, skip_connection_block_sets=[], 
@@ -613,58 +665,8 @@ def simplify_network_def_based_on_constraint(network_def, block, constraint, res
         Update # of input channels of one non-depthwise layer following the target layers.
     '''
     for current_num_out_channels in num_out_channels_try:  # Only allow multiple of '_MIN_FEATURE_SIZE'.
-        for target_layer_index in target_layer_indices:
-            update_num_out_channels = True
-            current_num_out_channels_after_pixel_shuffle = current_num_out_channels
-            for layer_idx, (layer_name, layer_properties) in enumerate(simplified_network_def.items()):
-                if layer_idx < target_layer_index:
-                    continue
-                
-                # for the block to be simplified (# of output channels is simplified)
-                if update_num_out_channels:
-                    if not layer_properties[KEY_IS_DEPTHWISE]:
-                        layer_properties[KEY_NUM_OUT_CHANNELS] = current_num_out_channels
-                        update_num_out_channels = False
-                        
-                        print('    simplify_def>     layer {}: num of output channel changed to {}'.format(layer_name, str(current_num_out_channels)))
-                    else:
-                        raise ValueError('Expected a non-depthwise layer but got a depthwise layer.')
-                # for blocks following the target blocks (# of input channels is simplified)
-                else:
-                    if current_num_out_channels_after_pixel_shuffle % layer_properties[
-                        KEY_BEFORE_SQUARED_PIXEL_SHUFFLE_FACTOR] != 0:
-                        raise ValueError('current_num_out_channels or current_num_out_channels_after_pixel_shuffle is '
-                                         'not divisible by the scaling factor of pixel shuffling.')
-                    current_num_out_channels_after_pixel_shuffle = (
-                            current_num_out_channels_after_pixel_shuffle / layer_properties[
-                        KEY_BEFORE_SQUARED_PIXEL_SHUFFLE_FACTOR])
-                    layer_properties[KEY_NUM_IN_CHANNELS] = current_num_out_channels_after_pixel_shuffle
-                    print('    simplify_def>     layer {}: num of input channel changed to {}'.format(layer_name, str(current_num_out_channels_after_pixel_shuffle)))
-
-                    '''
-                        Consider the case that a FC layer is placed after a Conv and Flatten:
-                            FC: input feature size: Cin
-                                output feature size: Cout
-                            Conv: output feature map size: H x W x C
-                            So Cin = H x W x C.
-                            If C -> C' based on constraints, then Cin -> H x W x C'
-                    '''
-                    if layer_properties[KEY_BEFORE_SQUARED_PIXEL_SHUFFLE_FACTOR] == 1:
-                        if network_def[layer_name][KEY_NUM_IN_CHANNELS] > max_num_out_channels:
-                            assert network_def[layer_name][KEY_NUM_IN_CHANNELS] % max_num_out_channels == 0
-                            # H x W here
-                            spatial_factor = network_def[layer_name][KEY_NUM_IN_CHANNELS] // max_num_out_channels
-                            layer_properties[KEY_NUM_IN_CHANNELS] = spatial_factor*current_num_out_channels
-                            print('    simplify_def>     [Update] layer {}: num of input channel changed to {}'.format(layer_name, str(spatial_factor*current_num_out_channels)))
-
-                    if not layer_properties[KEY_IS_DEPTHWISE]:
-                        break
-                    else:
-                        layer_properties[KEY_NUM_OUT_CHANNELS] = current_num_out_channels_after_pixel_shuffle
-                        layer_properties[KEY_GROUPS] = current_num_out_channels_after_pixel_shuffle
-                        print('    simplify_def>     depthwise layer {}: num of output channel changed to {}'.format(layer_name, str(current_num_out_channels_after_pixel_shuffle)))
-
-
+        # Get the simplified network def compatible with current_num_out_channels
+        simplified_network_def = compose_simplified_network_def(simplified_network_def, network_def, current_num_out_channels, target_layer_indices, max_num_out_channels)
         # Get the current resource consumption
         simplified_resource = compute_resource(simplified_network_def, resource_type, 
                                                lookup_table_path)
@@ -719,7 +721,7 @@ def simplify_model_based_on_network_def(simplified_network_def, model):
                     KEY_BEFORE_SQUARED_PIXEL_SHUFFLE_FACTOR]
                 kept_filter_idx = (kept_filter_idx[::before_squared_pixel_shuffle_factor] /
                                    before_squared_pixel_shuffle_factor)
-
+                kept_filter_idx = kept_filter_idx.long()
                 if layer_param_name == WEIGHTSTRING: #WEIGHTSTRING == layer_param_name:                    
                     if layer.groups == 1:  # Pointwise layer or depthwise layer with only one filter.
                         setattr(layer, layer_param_name,