FastAI
batchsize
(
batch
::
Union
{
Tuple
,
NamedTuple
}
)
=
batchsize
(
batch
[
1
]
)
batchsize
(
batch
::
Dict
)
=
batchsize
(
batch
[
first
(
keys
(
batch
)
)
]
)
batchsize
(
batch
::
AbstractArray
{
T
,
N
}
)
where
{
T
,
N
}
=
size
(
batch
,
N
)
unbatch
(
batch
)
=
collect
(
obsslices
(
batch
)
)
obsslices
(
batch
)
=
(
obsslice
(
batch
,
i
)
for
i
in
1
:
batchsize
(
batch
)
)
function
obsslice
(
batch
::
AbstractArray
{
T
,
N
}
,
i
)
where
{
T
,
N
}
return
view
(
batch
,
[
(
:
)
for
_
in
1
:
(
N
-
1
)
]
...
,
i
)
end
obsslice
(
batch
::
AbstractVector
,
i
)
=
batch
[
i
]
function
obsslice
(
batch
::
Tuple
,
i
)
return
Tuple
(
obsslice
(
batch
[
j
]
,
i
)
for
j
in
1
:
length
(
batch
)
)
end
function
obsslice
(
batch
::
NamedTuple
,
i
)
return
(
;
zip
(
keys
(
batch
)
,
obsslice
(
values
(
batch
)
,
i
)
)
...
)
end
function
obsslice
(
batch
::
Dict
,
i
)
return
Dict
(
k
=>
obsslice
(
v
,
i
)
for
(
k
,
v
)
in
batch
)
end