Документация Engee

Рекурсивные преобразования из Functors.jl

Модели Flux представляют собой структуры с глубоким вложением. Пакет Functors.jl предоставляет инструменты для анализа таких объектов, применения функций к содержащимся в них параметрам и их перестроения.

Новые слои должны аннотироваться с помощью макроса Functors.@functor. Это позволит функции params получать доступ к параметрам внутри модели, а функции gpu — перемещать их в GPU.

Для пакета Functors.jl имеются заметки о базовом использовании, в которых приведены дополнительные сведения. Кроме того, на странице Расширенное построение и настройка моделей подробнее рассматриваются сценарии использования Functors.

@functor T
@functor T (x,)

Добавляет в functor методы, позволяющие выполнять рекурсию по объектам типа T и их воссоздание. Предполагается, что конструктор типа T принимает все его поля. Это верно всегда, кроме случая, когда вы предоставили внутренний конструктор, не отвечающий этому условию.

По умолчанию все поля T считаются children. Набор таких полей можно ограничить путем предоставления кортежа имен полей.

Примеры

julia> struct Foo; x; y; end

julia> @functor Foo

julia> Functors.children(Foo(1,2))
(x = 1, y = 2)

julia> _, re = Functors.functor(Foo(1,2));

julia> re((10, 20))
Foo(10, 20)

julia> struct TwoThirds a; b; c; end

julia> @functor TwoThirds (a, c)

julia> ch2, re3 = Functors.functor(TwoThirds(10,20,30));

julia> ch2
(a = 10, c = 30)

julia> re3(("ten", "thirty"))
TwoThirds("ten", 20, "thirty")

julia> fmap(x -> 10x, TwoThirds(Foo(1,2), Foo(3,4), 56))
TwoThirds(Foo(10, 20), Foo(3, 4), 560)
fmap(f, x, ys...; exclude = Functors.isleaf, walk = Functors.DefaultWalk()[, prune])

Структура и тип с сохранением map.

По умолчанию преобразует каждый конечный узел (указываемый с помощью exclude, по умолчанию isleaf) путем применения f; в противном случае выполняет рекурсивный обход x с помощью functor. Эта функция может быть также связана с объектами ys с той же структурой дерева. В таком случае f применяется к соответствующим конечным узлам в x и ys.

Примеры

julia> fmap(string, (x=1, y=(2, 3)))
(x = "1", y = ("2", "3"))

julia> nt = (a = [1,2], b = [23, (45,), (x=6//7, y=())], c = [8,9]);

julia> fmap(println, nt)
[1, 2]
23
45
6//7
()
[8, 9]
(a = nothing, b = Any[nothing, (nothing,), (x = nothing, y = nothing)], c = nothing)

julia> fmap(println, nt; exclude = x -> x isa Array)
[1, 2]
Any[23, (45,), (x = 6//7, y = ())]
[8, 9]
(a = nothing, b = nothing, c = nothing)

julia> twice = [1, 2];  # println действует только один раз

julia> fmap(println, (i = twice, ii = 34, iii = [5, 6], iv = (twice, 34), v = 34.0))
[1, 2]
34
[5, 6]
34
34.0
(i = nothing, ii = nothing, iii = nothing, iv = (nothing, nothing), v = nothing)

julia> d1 = Dict("x" => [1,2], "y" => 3);

julia> d2 = Dict("x" => [4,5], "y" => 6, "z" => "an_extra_value");

julia> fmap(+, d1, d2) == Dict("x" => [5, 7], "y" => 9) # Обратите внимание, что z игнорируется
true

Изменяемые объекты, встречающиеся несколько раз, обрабатываются только один раз (путем кэширования f(x) в IdDict). Поэтому отношение x.i === x.iv[1] сохраняется. Неизменяемый объект, встречающийся дважды, не сохраняется в кэше, поэтому f(34) вызывается дважды и результаты согласуются, только если функция f чистая.

По умолчанию у Tuple, NamedTuple и некоторых других контейнерных типов в Base есть дочерние объекты для рекурсивного обхода. У массивов чисел их нет. Для обеспечения рекурсии по новым типам необходимо предоставить метод functor. Это можно сделать с помощью макроса @functor:

julia> struct Foo; x; y; end

julia> @functor Foo

julia> struct Bar; x; end

julia> @functor Bar

julia> m = Foo(Bar([1,2,3]), (4, 5, Bar(Foo(6, 7))));

julia> fmap(x -> 10x, m)
Foo(Bar([10, 20, 30]), (40, 50, Bar(Foo(60, 70))))

julia> fmap(string, m)
Foo(Bar("[1, 2, 3]"), ("4", "5", Bar(Foo("6", "7"))))

julia> fmap(string, m, exclude = v -> v isa Bar)
Foo("Bar([1, 2, 3])", (4, 5, "Bar(Foo(6, 7))"))

Для рекурсии по пользовательским типам без их последующего воссоздания используйте fmapstructure.

Для расширенной настройки обхода передайте пользовательскую функцию walk, являющуюся подтипом Functors.AbstractWalk. Вызов fmap(f, x, ys...; walk = mywalk) заключает mywalk в ExcludeWalk, а затем в CachedWalk. ExcludeWalk отвечает за применение f к исключаемым узлам. Сведения о низкоуровневом интерфейсе для выполнения пользовательского обхода см. в описании execute.

julia> struct MyWalk <: Functors.AbstractWalk end

julia> (::MyWalk)(recurse, x) = x isa Bar ? "hello" :
                                            Functors.DefaultWalk()(recurse, x)

julia> fmap(x -> 10x, m; walk = MyWalk())
Foo("hello", (40, 50, "hello"))

Поведение в ситуации, когда один и тот же узел встречается дважды, можно изменить, задав значение именованного аргумента prune, которое будет использоваться во всех случаях, кроме первого:

julia> twice = [1, 2];

julia> fmap(float, (x = twice, y = [1,2], z = twice); prune = missing)
(x = [1.0, 2.0], y = [1.0, 2.0], z = missing)
Functors.isleaf(x)

Возвращает true, если у x нет children в соответствии с functor.

Примеры

julia> Functors.isleaf(1)
true

julia> Functors.isleaf([2, 3, 4])
true

julia> Functors.isleaf(["five", [6, 7]])
false

julia> Functors.isleaf([])
false

julia> Functors.isleaf((8, 9))
false

julia> Functors.isleaf(())
true
Functors.children(x)

Возвращает дочерние объекты x согласно определению functor. Эквивалентно functor(x)[1].

fcollect(x; exclude = v -> false)

Выполняет обход x путем рекурсии по каждому дочернему объекту x согласно определению functor и собирает результаты в плоском массиве с упорядочением согласно обходу x в ширину и с учетом порядка итерации вызовов children.

Не выполняет рекурсию по ветвям, для которых корневыми являются узлы v с exclude(v) == true. В таких случаях корень v также исключается из результата. По умолчанию exclude всегда выдает false.

См. также описание children.

Примеры

julia> struct Foo; x; y; end

julia> @functor Foo

julia> struct Bar; x; end

julia> @functor Bar

julia> struct TypeWithNoChildren; x; y; end

julia> m = Foo(Bar([1,2,3]), TypeWithNoChildren(:a, :b))
Foo(Bar([1, 2, 3]), TypeWithNoChildren(:a, :b))

julia> fcollect(m)
4-element Vector{Any}:
 Foo(Bar([1, 2, 3]), TypeWithNoChildren(:a, :b))
 Bar([1, 2, 3])
 [1, 2, 3]
 TypeWithNoChildren(:a, :b)

julia> fcollect(m, exclude = v -> v isa Bar)
2-element Vector{Any}:
 Foo(Bar([1, 2, 3]), TypeWithNoChildren(:a, :b))
 TypeWithNoChildren(:a, :b)

julia> fcollect(m, exclude = v -> Functors.isleaf(v))
2-element Vector{Any}:
 Foo(Bar([1, 2, 3]), TypeWithNoChildren(:a, :b))
 Bar([1, 2, 3])
Functors.functor(x) = functor(typeof(x), x)

Возвращает кортеж, содержащий кортеж NamedTuple дочерних объектов x (обычно его поля) и функцию воссоздания. Управляет поведением fmap.

Для пользовательских типов следует добавить методы для functor(::Type{T}, x), обычно с помощью макроса @functor.

fmapstructure(f, x; exclude = isleaf)

Действует аналогично fmap, но не сохраняет тип пользовательских структур. Вместо этого возвращает NamedTuple, Tuple, массив или их вложенный набор.

Эта функция полезна в случае, когда результат не должен содержать пользовательских структур.

Примеры

julia> struct Foo; x; y; end

julia> @functor Foo

julia> m = Foo([1,2,3], [4, (5, 6), Foo(7, 8)]);

julia> fmapstructure(x -> 2x, m)
(x = [2, 4, 6], y = Any[8, (10, 12), (x = 14, y = 16)])

julia> fmapstructure(println, m)
[1, 2, 3]
4
5
6
7
8
(x = nothing, y = Any[nothing, (nothing, nothing), (x = nothing, y = nothing)])

Перемещение моделей или данных в GPU

Flux предоставляет ряд вспомогательных функций на основе fmap. Некоторые из них (f16, f32, f64) изменяют точность всех массивов в модели. Другие служат для перемещения модели в память GPU или из нее:

cpu(m)

Копирует m в память ЦП; функция, обратная gpu. Выполняет рекурсию по структурам, помеченным с помощью @functor.

Пример

julia> m_gpu = Dense(CUDA.randn(2, 5))
Dense(5 => 2)       # 12 параметров

julia> m_gpu.bias  # соответствует переданной матрице весов
2-element CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}:
 0.0
 0.0

julia> m = m_gpu |> cpu
Dense(5 => 2)       # 12 параметров

julia> m.bias
2-element Vector{Float32}:
 0.0
 0.0
gpu(m)

Копирует m на текущее устройство GPU (с использованием текущего бэкенда GPU), если оно доступно. Если устройство GPU недоступно, ничего не делает (но в первый раз выводит предупреждение).

Для массивов вызывает функцию cu CUDA, которая изменяет элементы Float64 массива на Float32 при их копировании на устройство (то же самое верно для AMDGPU). Для применения к массивам внутри структуры тип структуры должен быть помечен с помощью @functor.

Для копирования обратно в обычные массивы Array используйте функцию cpu. Сведения об изменении только типа элементов см. в описании функций f32 и f16.

В документации по CUDA.jl можно узнать, как определить текущее устройство.

Пример

julia> m = Dense(rand(2, 3))  # создается с матрицей весов типа Float64
Dense(3 => 2)       # 8 параметров

julia> typeof(m.weight)
Matrix{Float64} (alias for Array{Float64, 2})

julia> m_gpu = gpu(m)  # можно также записать в форме m_gpu = m |> gpu
Dense(3 => 2)       # 8 параметров

julia> typeof(m_gpu.weight)
CUDA.CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}
gpu(data::DataLoader)

Преобразует переданный объект DataLoader для применения функции gpu к каждому пакету данных при итерации по ним. (Если устройство GPU недоступно, ничего не делает.)

Пример

julia> dl = Flux.DataLoader((x = ones(2,10), y='a':'j'), batchsize=3)
4-element DataLoader(::NamedTuple{(:x, :y), Tuple{Matrix{Float64}, StepRange{Char, Int64}}}, batchsize=3)
  with first element:
  (; x = 2×3 Matrix{Float64}, y = 3-element StepRange{Char, Int64})

julia> first(dl)
(x = [1.0 1.0 1.0; 1.0 1.0 1.0], y = 'a':1:'c')

julia> c_dl = gpu(dl)
4-element DataLoader(::MLUtils.MappedData{:auto, typeof(gpu), NamedTuple{(:x, :y), Tuple{Matrix{Float64}, StepRange{Char, Int64}}}}, batchsize=3)
  with first element:
  (; x = 2×3 CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, y = 3-element StepRange{Char, Int64})

julia> first(c_dl).x
2×3 CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}:
 1.0  1.0  1.0
 1.0  1.0  1.0

Для больших наборов данных это предпочтительнее перемещения всех данных в память GPU перед созданием DataLoader:

julia> Flux.DataLoader((x = ones(2,10), y=2:11) |> gpu, batchsize=3)
4-element DataLoader(::NamedTuple{(:x, :y), Tuple{CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, UnitRange{Int64}}}, batchsize=3)
  with first element:
  (; x = 2×3 CUDA.CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, y = 3-element UnitRange{Int64})

Работает только при применении gpu непосредственно к DataLoader. Хотя функция gpu выполняется рекурсивно применительно к моделям Flux и многим базовым структурам Julia, она не будет работать, например, с кортежем DataLoader.