| 52 | return Architecture(genome) |
| 53 | |
| 54 | def build_model(self, architecture: Architecture, input_channels: int = 3, num_classes: int = 1000) -> nn.Module: |
| 55 | layers = [] |
| 56 | current_channels = input_channels |
| 57 | |
| 58 | for gene in architecture.genome: |
| 59 | if gene.get('type') == 'pooling': |
| 60 | if gene['pooling_type'] == 'max': |
| 61 | layers.append(nn.MaxPool2d(2)) |
| 62 | elif gene['pooling_type'] == 'avg': |
| 63 | layers.append(nn.AvgPool2d(2)) |
| 64 | else: |
| 65 | layer_type = gene['type'] |
| 66 | out_channels = gene['channels'] |
| 67 | activation = gene['activation'] |
| 68 | use_bn = gene['use_bn'] |
| 69 | dropout = gene['dropout'] |
| 70 | |
| 71 | if layer_type == 'conv3x3': |
| 72 | layers.append(nn.Conv2d(current_channels, out_channels, 3, padding=1)) |
| 73 | elif layer_type == 'conv5x5': |
| 74 | layers.append(nn.Conv2d(current_channels, out_channels, 5, padding=2)) |
| 75 | elif layer_type == 'conv7x7': |
| 76 | layers.append(nn.Conv2d(current_channels, out_channels, 7, padding=3)) |
| 77 | elif layer_type == 'depthwise': |
| 78 | layers.append(nn.Conv2d(current_channels, current_channels, 3, padding=1, groups=current_channels)) |
| 79 | layers.append(nn.Conv2d(current_channels, out_channels, 1)) |
| 80 | elif layer_type == 'bottleneck': |
| 81 | mid_channels = out_channels // 4 |
| 82 | layers.append(nn.Conv2d(current_channels, mid_channels, 1)) |
| 83 | if use_bn: |
| 84 | layers.append(nn.BatchNorm2d(mid_channels)) |
| 85 | layers.append(self._get_activation(activation)) |
| 86 | layers.append(nn.Conv2d(mid_channels, mid_channels, 3, padding=1)) |
| 87 | if use_bn: |
| 88 | layers.append(nn.BatchNorm2d(mid_channels)) |
| 89 | layers.append(self._get_activation(activation)) |
| 90 | layers.append(nn.Conv2d(mid_channels, out_channels, 1)) |
| 91 | elif layer_type == 'identity': |
| 92 | if current_channels != out_channels: |
| 93 | layers.append(nn.Conv2d(current_channels, out_channels, 1)) |
| 94 | else: |
| 95 | layers.append(nn.Identity()) |
| 96 | |
| 97 | if use_bn and layer_type != 'bottleneck': |
| 98 | layers.append(nn.BatchNorm2d(out_channels)) |
| 99 | |
| 100 | if layer_type != 'bottleneck': |
| 101 | layers.append(self._get_activation(activation)) |
| 102 | |
| 103 | if dropout > 0: |
| 104 | layers.append(nn.Dropout2d(dropout)) |
| 105 | |
| 106 | current_channels = out_channels |
| 107 | |
| 108 | layers.append(nn.AdaptiveAvgPool2d(1)) |
| 109 | layers.append(nn.Flatten()) |
| 110 | layers.append(nn.Linear(current_channels, num_classes)) |
| 111 | |