Rust in Everything
基于 burn 0.22.1

Config

主线位置：训练主线开始之前。模型的每一层都先写成一个 XxxConfig，再 init(device) 成模块；Learner 和优化器也用同样的配置类型（AdamConfig 等）。本页讲 #[derive(Config)] 生成了什么、配置怎么存取，以及为什么模块要分两段建；模块本身见 Module 派生。

flowchart TB
  cfg["LinearConfig new"] --> with["with_bias 等 setter"]
  with --> init["init device"]
  init --> module["Linear 模块"]
  cfg --> json["save 成 JSON"]
  json --> load["load 读回"]
  load --> init
  module --> rec["save_file 只存张量"]
  rec --> merge["load_file 到 init 出的模块"]

Config trait

Config 是 burn-core/src/config.rs 里一个很小的 trait：

pub trait Config: Debug + serde::Serialize + serde::de::DeserializeOwned {
    #[cfg(feature = "std")]
    fn save<P: AsRef<std::path::Path>>(&self, file: P) -> std::io::Result<()> {
        std::fs::write(file, config_to_json(self))
    }

    #[cfg(feature = "std")]
    fn load<P: AsRef<std::path::Path>>(file: P) -> Result<Self, ConfigError> {
        let content = std::fs::read_to_string(file.as_ref())
            .map_err(|_| ConfigError::FileNotFound(file.as_ref().to_string_lossy().to_string()))?;
        config_from_str(&content)
    }

    fn load_binary(data: &[u8]) -> Result<Self, ConfigError> {
        // ...
        config_from_str(content)
    }
}

三个方法都有默认实现，格式固定是 JSON：config_to_json 调 serde_json::to_string_pretty，config_from_str 调 serde_json::from_str，把错误转成 ConfigError::InvalidFormat。load 里任何读文件的失败（包括权限问题）都被报成 FileNotFound。trait 上没有 init：「XxxConfig::init 返回模块」是 burn-nn 里约定俗成的写法，不是 trait 方法。

#[derive(Config)] 生成什么

config_derive 进 derive_impl，按结构体或枚举选一个分析器，然后拼出六段代码：Config 的空实现、new、with_*、serde 两个 trait、Clone、Display。结构体的字段先被 create_struct_analyzer 分成三组：

组判定new 里setter
必填其余字段成为 new 的参数没有
可选类型名是 OptionNonewith_x(Option<T>)
有默认值写了 #[config(default = ...)]默认值with_x(T)

实验 crate module-exp 的 expand：

#[derive(Config, Debug)]
struct HeadConfig {
  d_model: usize,
  #[config(default = 1.0)]
  scale: f64,
  label: Option<String>,
}

cargo expand 的结果（删掉了生成的文档注释和 serde 自己展开出来的部分）：

impl burn::config::Config for HeadConfig {}
impl HeadConfig {
    #[allow(clippy::too_many_arguments)]
    pub fn new(d_model: usize) -> Self {
        Self {
            d_model,
            label: None,
            scale: 1.0,
        }
    }
}
impl HeadConfig {
    pub fn with_scale(mut self, scale: f64) -> Self {
        self.scale = scale;
        self
    }
    pub fn with_label(mut self, label: Option<String>) -> Self {
        self.label = label;
        self
    }
}
impl burn::serde::Serialize for HeadConfig {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: burn::serde::Serializer,
    {
        #[serde(crate = "burn::serde")]
        struct HeadConfigSerde {
            d_model: usize,
            label: Option<String>,
            scale: f64,
        }
        // ... serde 为 HeadConfigSerde 展开的 Serialize ...
        let serde_state = HeadConfigSerde {
            d_model: self.d_model.clone(),
            label: self.label.clone(),
            scale: self.scale.clone(),
        };
        serde_state.serialize(serializer)
    }
}
impl<'de> burn::serde::Deserialize<'de> for HeadConfig {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: burn::serde::Deserializer<'de>,
    {
        #[serde(crate = "burn::serde")]
        struct HeadConfigSerde {
            d_model: usize,
            label: Option<String>,
            scale: f64,
        }
        // ... serde 为 HeadConfigSerde 展开的 Deserialize ...
    }
}

（Clone 和 Display 没有贴：前者逐字段 clone，后者是 f.write_str(&burn::config::config_to_json(self))。）

几个细节都能从这段代码里读出来：

  • 字段顺序变了。 gen_new_fn 和 gen_serde_impl 按「必填、可选、默认」三组的顺序排字段，所以 JSON 里 label 排在 scale 前面，和声明顺序不同。
  • 序列化走一个影子结构体。 HeadConfigSerde 在函数体里现场声明、现场 #[derive(Serialize)]，serialize 时把每个字段 clone 进去再序列化。serde 派生怎么展开这个影子结构体，见 serde 卷的生成的 Serialize 和生成的 Deserialize。
  • 影子结构体上没有任何 #[serde(...)] 字段属性。 #[config(default = 1.0)] 只进了 new，没有变成 #[serde(default)]。
  • 默认值按字面量处理。 AttributeItem 只收字面量：数字、true 原样放进 new；字符串字面量会被 parse() 成一段 token 再放进去，所以 LinearConfig 能写 #[config(default = "Initializer::KaimingUniform{gain:1.0/num_traits::Float::sqrt(3.0), fan_out_only:false}")]，这段表达式在 new 被调用时才求值。
  • Debug 要自己派生，trait 要求它，宏不生成。生成的 impl HeadConfig { .. } 不带泛型参数，所以配置结构体不能是泛型的。

枚举（比如 Initializer、LinearLayout）也能 #[derive(Config)]，只生成 serde、Clone、Display 和 Config，没有 new 和 setter。

存取：JSON 和缺字段

实验 config（debug），把 HeadConfig 存成文件再读回，再喂几段手写的 JSON：

let config = HeadConfig::new(4).with_scale(2.0);
let path = dir.join("head.json");
config.save(&path).unwrap();
println!("saved:\n{}", std::fs::read_to_string(&path).unwrap());
println!("load back: {:?}", HeadConfig::load(&path));

for json in [
  r#"{"d_model": 4, "scale": 2.0}"#,
  r#"{"d_model": 4, "label": null}"#,
  r#"{"d_model": 4, "label": null, "scale": 2.0, "dropout": 0.1}"#,
] {
  println!("load_binary({json}) = {:?}", HeadConfig::load_binary(json.as_bytes()));
}
println!("missing file: {}", HeadConfig::load(dir.join("nope.json")).unwrap_err());

saved:
{
  "d_model": 4,
  "label": null,
  "scale": 2.0
}
load back: Ok(HeadConfig { d_model: 4, scale: 2.0, label: None })
load_binary({"d_model": 4, "scale": 2.0}) = Ok(HeadConfig { d_model: 4, scale: 2.0, label: None })
load_binary({"d_model": 4, "label": null}) = Err(InvalidFormat("missing field `scale` at line 1 column 29"))
load_binary({"d_model": 4, "label": null, "scale": 2.0, "dropout": 0.1}) = Ok(HeadConfig { d_model: 4, scale: 2.0, label: None })
missing file: Config error => File not found: out/cfg/nope.json
  • 缺 label（Option）可以：serde 派生对字面上是 Option 的字段缺省为 None。
  • 缺 scale（有默认值的字段）不行：影子结构体上没有 #[serde(default)]。这条规则和例外见 serde 卷的 missing field x。
  • 多出来的 dropout 被悄悄忽略：影子结构体没有 deny_unknown_fields。
📝NOTE
没有蠢问题

问： 字段上明明写了 #[config(default = 1.0)]，为什么 JSON 里缺了它还会报 missing field？

答： 这个默认值是给 new 用的，不是给反序列化用的。派生宏在 new 的函数体里写下 scale: 1.0，而反序列化走的是影子结构体 HeadConfigSerde，它的字段上一个 serde 属性都没有。后果是：给一个已经发布的配置加一个带默认值的新字段，之前存下来的 config.json 全都读不回来了。要兼容旧文件，只能把新字段声明成 Option<T>，在使用处自己补默认值。

LinearConfig 和 Conv2dConfig 打印出来（Display 就是 to_string_pretty）：

LinearConfig::new(4, 2) = {
  "d_input": 4,
  "d_output": 2,
  "bias": true,
  "initializer": {
    "KaimingUniform": {
      "gain": 0.5773502691896258,
      "fan_out_only": false
    }
  },
  "layout": "Col"
}
Conv2dConfig = {
  "channels": [
    3,
    8
  ],
  "kernel_size": [
    3,
    3
  ],
  ...
  "padding": "Valid",
  "bias": true,
  "initializer": {
    "KaimingUniform": {
      "gain": 0.5773502691896258,
      "fan_out_only": false
    }
  }
}

（前一个是 LinearConfig::new(4, 2).with_layout(LinearLayout::Col)。）初始化器也是配置的一部分，Initializer 是一个 #[derive(Config)] 的枚举，按 serde 的外部标签表示（见 serde 卷的外部标签）存成 {"KaimingUniform": {...}}。gain 是默认表达式 1.0 / sqrt(3.0) 算出来的值。

两段式：Config::init(device) → Module

burn-nn 的每个层都是一对类型：LinearConfig / Linear、Conv2dConfig / Conv2d、DropoutConfig / Dropout……LinearConfig::init 的签名是 pub fn init(&self, device: &Device) -> Linear。两种类型装的东西不同：

配置 LinearConfig模块 Linear
内容维度、开关、初始化器weight: Param<Tensor<2>>、bias: Option<Param<Tensor<1>>>
设备无参数所在的设备，init 时给定
序列化serde，JSON记录：只有张量（存储与加载）
能 deriveConfigModule

记录里只有参数张量：派生出来的 visit 根本不进常量字段（见 Module 派生），Flag 也被明确排除在记录之外。所以「步长是多少、有没有偏置、epsilon 是多少」这些信息只活在配置里，要重建一个模型，得先有配置。实验 config 的后半段：

let device = Device::flex();
let strided = Conv2dConfig::new([1, 1], [3, 3]).with_stride([2, 2]).init(&device);
let weight_sum: f32 = strided.weight.val().sum().into_scalar();
strided.save_file(dir.join("conv.bpk")).unwrap();
let loaded = Conv2dConfig::new([1, 1], [3, 3]).init(&device).load_file(dir.join("conv.bpk"));
let loaded_sum: f32 = loaded.weight.val().sum().into_scalar();
println!("saved stride [2, 2], weight sum {weight_sum}");
println!("loaded stride {:?}, weight sum {loaded_sum}", loaded.stride);
❗IMPORTANT
动脑筋：读回来的 loaded.stride 是多少？会报错吗？

存的是一个步长 [2, 2] 的卷积，读进一个用默认步长 [1, 1] 的配置 init 出来的卷积。两边权重形状都是 [1, 1, 3, 3]。

所以保存一个可以重建的模型，至少要两样东西：配置的 JSON（config.save）和参数的记录（model.save_file）。读的时候顺序固定：XxxConfig::load → init(&device) → load_file。init 之所以便宜，是因为参数是懒初始化的：init 只登记形状、设备和初始化闭包，紧接着的 load_file 用记录里的张量直接填上，随机初始化那一步从来没跑过。

init 要不要 device 由层决定：DropoutConfig::init() 不带参数，因为 Dropout 里只有一个概率和一个 Flag，没有任何张量。

动脑筋答案

saved stride [2, 2], weight sum -0.2514851
loaded stride [1, 1], weight sum -0.2514851

不报错，步长是 [1, 1]。权重和对得上，说明张量确实读进来了；步长是常量字段，不在记录里，所以保留了目标模块自己的值。这样读出来的卷积前向能跑，输出形状却和原模型不同，而且没有任何提示：记录里只有参数张量，步长没有东西可以拿来比对。

小结

  • Config trait 只有 save、load、load_binary 三个默认方法，格式固定是 pretty JSON；init 不是 trait 方法，是 burn-nn 的约定。
  • #[derive(Config)] 把字段分成必填、Option、有默认值三组，生成 new（只收必填）、with_*（后两组）、一个影子结构体做 serde、Clone 和打印 JSON 的 Display；Debug 要自己派生。
  • 默认值只在 new 里生效：反序列化时缺 Option 字段得到 None，缺有默认值的字段报 missing field，多余字段被忽略。
  • 模块分两段建，因为配置管超参数和形状、可序列化成 JSON，模块管参数张量、记录里只有张量；重建模型要配置和记录两样，读回顺序是 load → init → load_file，懒初始化让 init 不花代价。

💬 关联讨论

发起讨论
加载中...