How to generate the proper yolo style yaml?

[tobymuller233]

Search before asking

• I have searched the YOLOv5 issues and discussions and found no similar questions.

Question

Recently I'm working on something with yolo, and I have developed my own model and train it in this framework. When I tried to prune my model, I found that because of the particularity of the framework, everytime you need to train a model, you have to first have a yaml file telling the detailed structures of the network. Since pruned models have different channels, I have to modify the yaml file manually as discussed in this issue. So I'm now trying to design a script to automatically modify the yaml file for the network, though it can only be used for my network. I have tried a lot and failed to dump the correct format that looks same with original yaml format. Here is my code:

import yaml
import torch
import argparse
from models.common import *
from models.yolo import *
from ruamel.yaml import YAML
from ruamel.yaml.comments import CommentedSeq

def make_compact_list(data):
    result = CommentedSeq(data)
    result.fa.set_flow_style()  
    return result

def parse_model(model_dict, model):
# parse the model and generate yaml dict

    for i in range(len(model.model)):
        if i < 22:
            part = "backbone"
        else:
            part = "head"
        if isinstance(model.model[i], Conv):
            model_dict[part].append([-1, 1, "Conv", [model.model[i].conv.out_channels, model.model[i].conv.kernel_size[0], model.model[i].conv.stride[0]]])
        elif isinstance(model.model[i], DWConv):
            model_dict[part].append([-1, 1, "DWConv", [model.model[i].conv.out_channels, model.model[i].conv.kernel_size[0], model.model[i].conv.stride[0]]])
        elif isinstance(model.model[i], Bottleneck3):
            model_dict[part].append([-1, 1, "Bottleneck3", [model.model[i].cv3.conv.out_channels, model.model[i].cv1.conv.out_channels]])
        elif isinstance(model.model[i], nn.Sequential):
            for j in range(len(model.model[i])):
                # all Bottleneck3
                model_dict[part].append([-1, 1, "Bottleneck3", [model.model[i][j].cv3.conv.out_channels, model.model[i][j].cv1.conv.out_channels]])
        elif isinstance(model.model[i], Concat):
            if i == 23:
                model_dict[part].append([[-1, -5], 1, "Concat", [1]])
            elif i == 29:
                model_dict[part].append([[-1, 12], 1, "Concat", [1]])
            elif i == 35:
                model_dict[part].append([[-1, 7], 1, "Concat", [1]])
            else:
                # error
                print(f"Error: Concat layer position ({i}) is wrong")
        elif isinstance(model.model[i], nn.Upsample):
            model_dict[part].append([-5, 1, "nn.Upsample", [None, 2, "nearest"]])
        elif isinstance(model.model[i], Detect):
            model_dict[part].append([[44, 38, 32], 1, "Detect", ["nc", "anchors"]])
        else:
            # error
            print(f"Error: Layer type is not supported: {model.model[i]}")
    model_dict['backbone'] = make_compact_list(model_dict['backbone'])
    model_dict['head'] = make_compact_list(model_dict['head'])
    
if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Generate yaml file for the pruned model")
    parser.add_argument("--model", type=str, help="Path to the pruned model")

    opt = parser.parse_args()
    model = torch.load(opt.model)['model']          # Load the pruned model
    print(model)
    # Create the yaml file
    model_dict = {}
    model_dict['nc'] = 1
    model_dict['depth_multiple'] = 1.0
    model_dict['width_multiple'] = 1.0
    anchors = [[4, 6, 7, 10, 11, 15], [16, 24, 33, 25, 26, 41], [47, 60, 83, 97, 141, 149]]
    model_dict['anchors'] = make_compact_list(anchors)
    model_dict['backbone'] = []
    model_dict['head'] = []
    
    parse_model(model_dict, model)
    yaml_save = YAML()
    
    yaml.PreserveAnchor = False
    yaml_save.default_block_style = True
    yaml_save.indent(sequence=4, offset=2)
    with open("pruned_model.yaml", "w") as f:
        yaml_save.dump(model_dict, f)

Here shows what it looks like:

nc: 1
depth_multiple: 1.0
width_multiple: 1.0
anchors: [[4, 6, 7, 10, 11, 15], [16, 24, 33, 25, 26, 41], [47, 60, 83, 97, 141, 149]]
backbone: [[-1, 1, Conv, [2, 3, 2]], [-1, 1, Conv, [2, 3, 1]], [-1, 1, Conv, [7, 1,
          1]], [-1, 1, Conv, [19, 1, 1]], [-1, 1, Conv, [6, 1, 1]], [-1, 1, Bottleneck3,
      [6, 34]], [-1, 1, Conv, [32, 1, 1]], [-1, 1, Conv, [32, 3, 2]], [-1, 1, Conv,
      [8, 1, 1]], [-1, 1, Bottleneck3, [8, 33]], [-1, 1, Bottleneck3, [8, 44]], [
      -1, 1, Conv, [38, 1, 1]], [-1, 1, Conv, [38, 3, 2]], [-1, 1, Conv, [12, 1, 1]],
   [-1, 1, Bottleneck3, [12, 78]], [-1, 1, Bottleneck3, [12, 89]], [-1, 1, Bottleneck3,
      [12, 88]], [-1, 1, Conv, [83, 1, 1]], [-1, 1, Conv, [83, 3, 1]], [-1, 1, Conv,
      [24, 1, 1]], [-1, 1, Bottleneck3, [24, 113]], [-1, 1, Bottleneck3, [24, 132]],
   [-1, 1, Conv, [115, 1, 1]], [-1, 1, Conv, [115, 3, 2]], [-1, 1, Conv, [25, 1, 1]],
   [-1, 1, Bottleneck3, [25, 130]], [-1, 1, Bottleneck3, [25, 218]]]
head: [[-1, 1, Conv, [64, 1, 1]], [[-1, -5], 1, Concat, [1]], [-1, 1, Conv, [39, 1,
          1]], [-1, 1, Conv, [39, 3, 1]], [-1, 1, Conv, [33, 1, 1]], [-1, 1, Conv,
      [18, 1, 1]], [-5, 1, nn.Upsample, [null, 2, nearest]], [[-1, 12], 1, Concat,
      [1]], [-1, 1, Conv, [19, 1, 1]], [-1, 1, Conv, [19, 3, 1]], [-1, 1, Conv, [
          21, 1, 1]], [-1, 1, Conv, [18, 1, 1]], [-5, 1, nn.Upsample, [null, 2, nearest]],
   [ [-1, 7], 1, Concat, [1]], [-1, 1, Conv, [13, 1, 1]], [-1, 1, Conv, [13, 3, 1]],
   [-1, 1, Conv, [16, 1, 1]], [-1, 1, Conv, [18, 1, 1]], [[44, 38, 32], 1, Detect,
      [nc, anchors]]]

And I just want the members of backbone and head be in a single row just like this:

nc: 1 # number of classes
depth_multiple: 1.0 # model depth multiple
width_multiple: 1.0 # layer channel multiple
anchors:
  - [4, 6, 7, 10, 11, 15]
  - [16, 24, 33, 25, 26, 41]
  - [47, 60, 83, 97, 141, 149]

backbone:
  # [from, number, module, args]
  # args: out_channels, size, stride
  [
    [-1, 1, Conv, [2, 3, 2]],  # 0  [batch, 8, size/2, size/2]
    [-1, 1, DWConv, [2, 3, 1]], # 1 [320]
    [-1, 1, Conv, [7, 1, 1 ]], # 2  [320]
    [-1, 1, Conv, [19, 1, 1]], # 3 [-1, 1, DWConv, [24, 3, 2]] # 4
    [-1, 1, Conv, [6, 1, 1]], # 4
    [-1, 1, Bottleneck3, [6, 34]], # 5

    [-1, 1, Conv, [32, 1, 1]], # 6  
    [-1, 1, DWConv, [32, 3, 2]], # 7  [160]
    [-1, 1, Conv, [8, 1, 1]], # 8
    [-1, 1, Bottleneck3, [8, 33]], # 9
    [-1, 1, Bottleneck3, [8, 44]], # 10
    
    [-1, 1, Conv, [38, 1, 1]], # 11 
    [-1, 1, DWConv, [38, 3, 2]], # 12 [80] 
    [-1, 1, Conv, [12, 1, 1]], # 13
    [-1, 1, Bottleneck3, [12, 78]], # 14
    [-1, 1, Bottleneck3, [12, 89]], # 15
    [-1, 1, Bottleneck3, [12, 88]], # 16

    [-1, 1, Conv, [83, 1, 1]], # 17
    [-1, 1, DWConv, [83, 3, 1]], # 18
    [-1, 1, Conv, [24, 1, 1]], # 19
    [-1, 1, Bottleneck3, [24, 113]], # 20
    [-1, 1, Bottleneck3, [24, 132]], # 21

    [-1, 1, Conv, [115, 1, 1]], # 22    [80]
    [-1, 1, DWConv, [115, 3, 2]], # 23  [80] -> [40] 
    [-1, 1, Conv, [25, 1, 1]], # 24
    [-1, 1, Bottleneck3, [25, 130]], # 25 [batch, 40, size/16, size/16]
    [-1, 1, Bottleneck3, [25, 218]], # 26 [batch, 40, size/16, size/16]
  ]

head: [
    [-1, 1, Conv, [64, 1, 1]], # 27 [40]
    [[-1, -5], 1, Concat, [1]], # 28  [batch, 224, size/16, size/16]  [40]  # to line 40 # changed from -4 to -5

    [-1, 1, Conv, [39, 1, 1]], # 29
    [-1, 1, DWConv, [39, 3, 1]], # 30
    [-1, 1, Conv, [33, 1, 1]], # 31
    [-1, 1, Conv, [18, 1, 1]], # 32   [batch, 18, size/8, size/8] -> [40] ###
    
    [-5, 1, nn.Upsample, [None, 2, "nearest"]],  # 33   [80]
    [[-1, 12], 1, Concat, [1]],  # 34   [80]  ch = 272      # to line 27  # changed from 11 to 12
    [-1, 1, Conv, [19, 1, 1]], # 35
    [-1, 1, DWConv, [19, 3, 1]], # 36 
    [-1, 1, Conv, [21, 1, 1]], # 37   
    [-1, 1, Conv, [18, 1, 1]], # 38 [batch, 18, 160, 160] -> [80] ###

    [-5, 1, nn.Upsample, [None, 2, "nearest"]],  # 39 [1, 272, 320, 320] -> [160]
    [[-1, 7], 1, Concat, [1]],  # 40  # to line 21
    [-1, 1, Conv, [13, 1, 1]], # 41   
    [-1, 1, DWConv, [13, 3, 1]], # 42 
    [-1, 1, Conv, [16, 1, 1]], # 43   
    [-1, 1, Conv, [18, 1, 1]], # 44   [batch, 18, 320, 320] -> [160]  ###

    [[44, 38, 32], 1, Detect, [nc, anchors]], 
  ]

FYI, the reason for why I used ruamel.yaml instead of yaml is that I tried yaml before using:

def dump_yaml(data, file_path):
    class MyDumper(yaml.Dumper):
        def increase_indent(self, flow=False, indentless=False):
            return super(MyDumper, self).increase_indent(flow=flow, indentless=indentless)
    with open(file_path, 'w') as f:
        yaml.dump(data, f, Dumper=MyDumper, default_flow_style=None)

But it turns out that this only fits anchors since there is only one layer of embedded list instead of two.:

anchors:
- [4, 6, 7, 10, 11, 15]
- [16, 24, 33, 25, 26, 41]
- [47, 60, 83, 97, 141, 149]
backbone:
- - -1
  - 1
  - Conv
  - [2, 3, 2]
- - -1
  - 1
  - Conv
  - [2, 3, 1]
...

Additional

No response

[UltralyticsAssistant]

👋 Hello @tobymuller233, thank you for your interest in YOLOv5 🚀! It sounds like you're diving deep into modifying the YAML configuration for custom network architectures and pruning.

If this is a 🐛 Bug Report, please ensure you include a minimum reproducible example to help us assist you more effectively.

For custom YAML modification ❓ Questions, please provide details about your dataset and training process, and make sure you’ve followed best practices. It's fantastic that you're working on automating this process; remember that ruamel.yaml can be very handy for preserving YAML structure while editing.

Requirements

Ensure you are using Python>=3.8.0 with all dependencies from requirements.txt installed, including PyTorch>=1.8.

Environments

YOLOv5 can run in various environments with the necessary dependencies installed like Notebooks with GPU support, Google Cloud, AWS, and Docker containers.

Status

If you encounter issues with YAML formatting, note the indentations and flow styles of the YAML libraries you are using. This can affect how nested lists are displayed, as you've observed.

This is an automated response, and an Ultralytics engineer will be with you shortly to provide further assistance. Thank you for your patience and for sharing your work with us! 😊

[pderrenger]

To generate a YOLO-style YAML file with specific formatting, you can use the ruamel.yaml library, which allows for more control over how lists and other structures are formatted. In your script, ensure you set the flow style correctly for nested lists. Here's a refined approach to achieve your desired format:

from ruamel.yaml import YAML
from ruamel.yaml.comments import CommentedSeq

def make_compact_list(data):
    result = CommentedSeq(data)
    result.fa.set_flow_style()  # Set flow style for inline list representation
    return result

# ... (rest of your existing code)

yaml_save = YAML()
yaml_save.default_flow_style = None  # Ensures proper inline formatting
yaml_save.indent(sequence=4, offset=2)  # Adjust indentation

with open("pruned_model.yaml", "w") as f:
    yaml_save.dump(model_dict, f)

Make sure to apply make_compact_list to both backbone and head lists within your parse_model function. This should help in maintaining the inline format you are aiming for. If further issues persist, consider checking if ruamel.yaml is appropriately handling nested structures as expected.