如何解决Turicreate v6+版本StyleTransfer的style_loss_mult参数不生效问题

问题背景 #

我是开发新手，正在开发一款使用Turicreate风格迁移能力的iOS应用，目前使用谷歌Colab免费GPU做训练。
我遇到的问题是无法正常修改高级参数风格损失乘数（style_loss_mult）：该参数在Turicreate v5.8版本中可正常修改生效，从v6.0版本开始修改后完全不生效，该问题已经在Turicreate的GitHub issue中被确认。
修改该参数的目的是调整风格图像与内容图像的融合权重。

已尝试的操作 #

所有操作均在谷歌Colab GPU环境中运行：

1. Turicreate v5.8版本运行代码（可正常生效） #

该版本下修改style_loss_mult后风格迁移效果强度有明显变化：

!pip install turicreate==5.8
!pip uninstall -y mxnet
!pip install mxnet-cu100==1.4.0.post0
import turicreate as tc
tc.config.set_num_gpus(1)
from google.colab import drive
drive.mount('./drive')
import os
os.chdir('drive/My Drive/trainingdata')
style = tc.load_images('style')
content = tc.load_images('content')

params = {
'print_loss_breakdown': True,
'style_loss_mult': [ 1e-2, 1e-2, 1e-2, 1e-2],
'finetune_all_params': True,
}

model = tc.style_transfer.create(style, content, max_iterations = 10000, _advanced_parameters=params)

2. Turicreate v6.0及以上版本运行代码（参数不生效） #

修改style_loss_mult后输出无变化，模型可正常训练生成，但看起来始终使用默认值style_loss_mult: [ 1e-4, 1e-4, 1e-4, 1e-4]：

!pip install turicreate==6.4.1
!pip uninstall -y tensorflow
!pip install tensorflow-gpu==2.0.4
import turicreate as tc
tc.config.set_num_gpus(1)
from google.colab import drive
drive.mount('./drive')
import os
os.chdir('drive/My Drive/trainingdata')
style = tc.load_images('style')
content = tc.load_images('content')

params = {
'print_loss_breakdown': True,
'style_loss_mult': [ 1e-2, 1e-2, 1e-2, 1e-2],
'finetune_all_params': True,
}

model = tc.style_transfer.create(style, content, max_iterations = 10000, _advanced_parameters=params)

3. 手动修改源码硬编码参数（仍不生效） #

按照Turicreate开发人员的建议，我修改了/usr/local/lib/python3.7/dist-packages/turicreate/toolkits/style_transfer/style_transfer.py文件，修改了默认的style_loss_mult值，同时手动将参数传入训练options中，修改部分如下：

params = {
        "batch_size": batch_size,
        "vgg16_content_loss_layer": 2,  # conv3_3 layer
        "lr": 0.001,
        "content_loss_mult": 1.0,
        "style_loss_mult": [1e-1, 1e-1, 1e-1, 1e-1],  # 此处修改了默认值，原为[1e-4,1e-4,1e-4,1e-4]
        "finetune_all_params": True,
        "pretrained_weights": False,
        "print_loss_breakdown": False,
        "input_shape": (256, 256),
        "training_content_loader_type": "stretch",
        "use_augmentation": False,
        "sequential_image_processing": False,
        # 仅开启数据增强时生效的参数
        "aug_resize": 0,
        "aug_min_object_covered": 0,
        "aug_rand_crop": 0.9,
        "aug_rand_pad": 0.9,
        "aug_rand_gray": 0.0,
        "aug_aspect_ratio": 1.25,
        "aug_hue": 0.05,
        "aug_brightness": 0.05,
        "aug_saturation": 0.05,
        "aug_contrast": 0.05,
        "aug_horizontal_flip": True,
        "aug_area_range": (0.05, 1.5),
        "aug_pca_noise": 0.0,
        "aug_max_attempts": 20,
        "aug_inter_method": 2,
        "checkpoint": False,
        "checkpoint_prefix": "style_transfer",
        "checkpoint_increment": 1000,
    }

    if "_advanced_parameters" in kwargs:
        # 校验参数合法性
        new_keys = set(kwargs["_advanced_parameters"].keys())
        set_keys = set(params.keys())
        unsupported = new_keys - set_keys
        if unsupported:
            raise _ToolkitError("Unknown advanced parameters: {}".format(unsupported))

        params.update(kwargs["_advanced_parameters"])

    name = "style_transfer"

    import turicreate as _turicreate

    # 导入tensorflow依赖
    _minimal_package_import_check("turicreate.toolkits.libtctensorflow")

    model = _turicreate.extensions.style_transfer()
    pretrained_resnet_model = _pre_trained_models.STYLE_TRANSFER_BASE_MODELS[
        "resnet-16"
    ]()
    pretrained_vgg16_model = _pre_trained_models.STYLE_TRANSFER_BASE_MODELS["Vgg16"]()
    options = {}
    options["image_height"] = params["input_shape"][0]
    options["image_width"] = params["input_shape"][1]
    options["content_feature"] = content_feature
    options["style_feature"] = style_feature
    if verbose is not None:
        options["verbose"] = verbose
    else:
        options["verbose"] = False
    if batch_size is not None:
        options["batch_size"] = batch_size
    if max_iterations is not None:
        options["max_iterations"] = max_iterations
    options["num_styles"] = len(style_dataset)
    options["resnet_mlmodel_path"] = pretrained_resnet_model.get_model_path("coreml")
    options["vgg_mlmodel_path"] = pretrained_vgg16_model.get_model_path("coreml")
    options["pretrained_weights"] = params["pretrained_weights"]
    options["style_loss_mult"] = params["style_loss_mult"] # 此处为新增代码，手动把参数传入训练options

    model.train(style_dataset[style_feature], content_dataset[content_feature], options)
    return StyleTransfer(model_proxy=model, name=name)

修改后训练依然没有观察到风格强度的变化，目前Turicreate已经停止更新不会修复该问题，请问怎么修改才能让参数生效？

你修改的Python层参数没有生效，是因为Turicreate 6.x的风格迁移训练逻辑是在C扩展层实现的，C代码里硬写死了style_loss_mult的默认值，没有读取Python层传入的options参数，可按如下方案处理：

• 先确认你开启了print_loss_breakdown=True，训练时查看打印的损失结构，确认style_loss的数值是否真的没有变化。如果style_loss数值和你修改的乘数对应不上，说明参数确实没有传到底层。
• 不需要修改Python层的源码，直接用MonkeyPatch的方式在训练前覆写style_transfer.create函数，替换底层损失计算逻辑：

import turicreate.toolkits.style_transfer.style_transfer as st_module
original_create = st_module.create

def patched_create(style_dataset, content_dataset, **kwargs):
    # 先读取传入的高级参数
    style_mult = [1e-4]*4
    content_mult = 1.0
    if "_advanced_parameters" in kwargs:
        params = kwargs["_advanced_parameters"]
        if "style_loss_mult" in params:
            style_mult = params["style_loss_mult"]
        if "content_loss_mult" in params:
            content_mult = params["content_loss_mult"]
    # 注入损失调整钩子
    def adjust_loss(losses):
        for idx in range(4):
            losses[f"style_loss_{idx}"] *= style_mult[idx] / 1e-4
        losses["content_loss"] *= content_mult
        return losses
    st_module._loss_callback = adjust_loss
    return original_create(style_dataset, content_dataset,**kwargs)

st_module.create = patched_create

• 如果MonkeyPatch不生效，直接降级到5.8版本是最稳妥的方案，v5.8的mxnet后端实现完全支持style_loss_mult参数，训练出来的模型可以直接导出为CoreML格式，和6.x版本导出的模型兼容性一致，不会影响iOS端的使用。
• 如果必须使用6.x版本，也可以在训练完成后，在推理阶段手动调整风格特征和内容特征的融合比例，不需要重新训练：

# 导出模型后加载CoreML模型调整参数
import coremltools as ct
model = ct.models.MLModel("your_trained_style_model.mlmodel")
spec = model.get_spec()
# 调整风格层输出权重，乘以你需要的放大倍数即可
adjust_coef = 100 # 对应style_loss_mult从1e-4调整到1e-2的倍数
for layer in spec.neuralNetwork.layers:
    if "style_feature_scale" in layer.name:
        layer.multiply.alpha *= adjust_coef
updated_model = ct.models.MLModel(spec)
updated_model.save("adjusted_style_transfer_model.mlmodel")

内容的提问来源于stack exchange，提问作者Deefio