R中brms函数内保存模型时RDS文件体积异常增大的解决方法

我在R中使用brms包拟合逐词阅读时长实验的模型，为批量处理不同词汇数据，将拟合代码封装为函数，并通过RDS文件保存模型避免重复拟合。但发现函数内保存的RDS文件体积随模型数量递增：第一个含11个参数的模型约141MB，第二个同参数模型约282MB，第三个约423MB。最小可复现代码如下：

library(brms)

fit.models <- function() {
    set.seed(0)
    m1 <- brm(
        formula = Sepal.Length ~ Sepal.Width,
        data = iris,
        cores = 4,
        chains = 4,
        iter = 100,
        file = 'm1-function.rds'
    )
    
    set.seed(0)
    m2 <- brm(
        formula = Sepal.Length ~ Petal.Length,
        data = iris,
        cores = 4,
        chains = 4,
        iter = 100,
        file = 'm2-function.rds'
    )
}

fit.models()

set.seed(0)
m1 <- brm(
    formula = Sepal.Length ~ Sepal.Width,
    data = iris,
    cores = 4,
    chains = 4,
    iter = 100,
    file = 'm1-global.rds'
)

set.seed(0)
m2 <- brm(
    formula = Sepal.Length ~ Petal.Length,
    data = iris,
    cores = 4,
    chains = 4,
    iter = 100,
    file = 'm2-global.rds'
)

运行后m2-function.rds体积约为m1-function.rds的两倍，而全局环境保存的两个RDS体积相近。推测是brms对调用环境的处理导致冗余数据被存入RDS。

原因分析 #

brms生成的brmsfit对象会保留其调用环境（即拟合时的上下文环境）。在函数内连续拟合多个模型时，后续模型的调用环境会包含函数内已创建的所有对象（比如m1会留在函数环境中，被m2的环境引用），导致保存m2时，连带把m1的完整内容也打包进了RDS文件，造成体积翻倍。而全局环境中每个模型的环境相互独立，不会互相引用，因此体积正常。

解决方案 #

以下三种方法均可解决该问题：

方法1：拟合后立即移除前序模型对象 #

在函数内每次拟合完一个模型后，用rm()移除该对象，清理函数环境，避免后续模型的环境引用冗余内容：

fit.models <- function() {
    set.seed(0)
    m1 <- brm(
        formula = Sepal.Length ~ Sepal.Width,
        data = iris,
        cores = 4,
        chains = 4,
        iter = 100,
        file = 'm1-function.rds'
    )
    rm(m1) # 移除已保存的模型，清理函数环境
    
    set.seed(0)
    m2 <- brm(
        formula = Sepal.Length ~ Petal.Length,
        data = iris,
        cores = 4,
        chains = 4,
        iter = 100,
        file = 'm2-function.rds'
    )
}

方法2：重置模型的调用环境 #

手动将模型的调用环境修改为全局环境或空环境，避免携带函数内的上下文数据：

fit.models <- function() {
    set.seed(0)
    m1 <- brm(
        formula = Sepal.Length ~ Sepal.Width,
        data = iris,
        cores = 4,
        chains = 4,
        iter = 100
    )
    environment(m1$call) <- globalenv() # 将模型的调用环境改为全局环境
    saveRDS(m1, 'm1-function.rds')
    rm(m1)
    
    set.seed(0)
    m2 <- brm(
        formula = Sepal.Length ~ Petal.Length,
        data = iris,
        cores = 4,
        chains = 4,
        iter = 100
    )
    environment(m2$call) <- globalenv()
    saveRDS(m2, 'm2-function.rds')
}

注意：此处不要使用brm的file参数，需手动修改环境后再用saveRDS保存。

方法3：用局部环境隔离每个模型的拟合过程 #

使用local()为每个模型创建独立的临时环境，确保不同模型的拟合上下文互不干扰：

fit.models <- function() {
    # 在独立局部环境中拟合第一个模型
    local({
        set.seed(0)
        m1 <- brm(
            formula = Sepal.Length ~ Sepal.Width,
            data = iris,
            cores = 4,
            chains = 4,
            iter = 100,
            file = 'm1-function.rds'
        )
    })
    
    # 在独立局部环境中拟合第二个模型
    local({
        set.seed(0)
        m2 <- brm(
            formula = Sepal.Length ~ Petal.Length,
            data = iris,
            cores = 4,
            chains = 4,
            iter = 100,
            file = 'm2-function.rds'
        )
    })
}

验证 #

使用上述任意方法后，函数内保存的RDS文件体积会与全局环境保存的文件体积相近，不会出现随模型数量递增的情况。

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