frozenleaves/accelerate
1
0
Fork 0
mirror of https://github.com/huggingface/accelerate.git synced 2025-10-20 10:03:46 +08:00
Code Issues Packages Projects Releases Wiki Activity

Add source code for DataLoader Animation (#2696)

Browse Source 
* dl animation

* oops

* Export
This commit is contained in:
Zach Mueller
2024-04-23 04:28:28 -04:00
committed by GitHub
parent ef0f62c12a
commit 3dc131cd8d
8 changed files with 903 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

32
manim_animations/dataloaders/stage_0.py Normal file
View File

@ -0,0 +1,32 @@
# Copyright 2024 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from manim import *
class Stage0(Scene):
def construct(self):
mascot = ImageMobject("mascot_bookie.png")
mascot.scale(.35)
mascot.move_to([-3.75,-1,0])
text = Paragraph(
"Distributed Training,\nHugging Face Accelerate,\nand PyTorch DataLoaders\n\nHow do they all interact?",
font_size=36,
line_spacing=1,
alignment="center",
weight=BOLD,
)
text.move_to([1.75,.5,0])
self.add(mascot)
self.add(text)

31
manim_animations/dataloaders/stage_1.py Normal file
View File

@ -0,0 +1,31 @@
# Copyright 2024 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from manim import *
class Stage01(Scene):
def construct(self):
mascot = ImageMobject("mascot_bookie.png")
mascot.scale(.35)
mascot.move_to([-3.75,-1,0])
text = Paragraph(
"Distributed Training,\nHugging Face Accelerate,\nand PyTorch DataLoaders\n\nHow do they all interact?",
font_size=36,
line_spacing=1,
alignment="center",
weight=BOLD,
)
text.move_to([1.75,.5,0])
self.add(mascot)
self.add(text)

176
manim_animations/dataloaders/stage_2.py Normal file
View File

@ -0,0 +1,176 @@
# Copyright 2024 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from manim import *
class Stage2(Scene):
def construct(self):
# The dataset items
fill = Rectangle(height=0.46,width=0.46).set_stroke(width=0)
columns = [
VGroup(*[Rectangle(height=0.25,width=0.25,color="green") for i in range(8)]).arrange(RIGHT,buff=0)
for j in range(4)
]
dataset_recs = VGroup(*columns).arrange(UP, buff=0)
dataset_text = Text("Dataset", font_size=24)
dataset = Group(dataset_recs,dataset_text).arrange(DOWN, buff=0.5, aligned_edge=DOWN)
dataset.move_to([-2,0,0])
self.add(dataset)
code = Code(
code="dataloader = DataLoader(...)\nfor batch in dataloader():\n\t...",
tab_width=4,
background="window",
language="Python",
font="Monospace",
font_size=14,
corner_radius=.2,
insert_line_no=False,
line_spacing=.75,
style=Code.styles_list[1],
)
code.move_to([-3.5, 2.5, 0])
self.add(code)
# The dataloader itself
dataloader = Group(
Rectangle(color="red", height=2, width=2),
Text("DataLoader", font_size=24)
).arrange(DOWN, buff=.5, aligned_edge=DOWN)
sampler = Group(
Rectangle(color="blue", height=1, width=1),
Text("Sampler", font_size=12)
).arrange(DOWN, buff=.25, aligned_edge=DOWN)
dataloader.move_to([1, 0, 0])
sampler.move_to([.75,.25,0])
self.add(dataloader)
self.add(sampler)
gpu_1 = Group(
Rectangle(color="white", height=1, width=1),
Text("GPU 1", font_size=12)
).arrange(DOWN, buff=.25, aligned_edge=DOWN).move_to([4, 2, 0])
gpu_2 = Group(
Rectangle(color="white", height=1, width=1),
Text("GPU 2", font_size=12)
).arrange(DOWN, buff=.25, aligned_edge=DOWN).move_to([4, .5, 0])
gpu_3 = Group(
Rectangle(color="white", height=1, width=1),
Text("GPU 3", font_size=12)
).arrange(DOWN, buff=.25, aligned_edge=DOWN).move_to([4, -1, 0])
gpu_4 = Group(
Rectangle(color="white", height=1, width=1),
Text("GPU 4", font_size=12)
).arrange(DOWN, buff=.25, aligned_edge=DOWN).move_to([4, -2.5, 0])
gpus = [gpu_1[0], gpu_2[0], gpu_3[0], gpu_4[0]]
self.add(gpu_1, gpu_2, gpu_3, gpu_4)
# Animate their existence
self.play(
Create(gpu_1[0], run_time=0.5),
Create(gpu_2[0], run_time=0.5),
Create(gpu_3[0], run_time=0.5),
Create(gpu_4[0], run_time=0.5),
Create(dataset_recs, run_time=1),
Create(sampler[0], run_time=1),
Create(dataloader[0], run_time=1)
)
step_1 = MarkupText(
f"Without any special care, \nthe same data is sent though each sampler, \nand the same samples are spit out on each GPU",
font_size=18
)
step_1.move_to([0, -2.5, 0])
self.play(
Write(step_1, run_time=4),
)
first_animations = []
second_animations = []
colors = ["BLUE_E", "DARK_BROWN", "GOLD_E", "GRAY_A"]
current_color = colors[0]
buff = 0
lr_buff = .25
old_target = None
new_datasets = []
for i,data in enumerate(dataset_recs[-1]):
if i % 2 == 0:
# current_color = colors[i//2]
current_color = "BLUE_E"
dataset_target = Rectangle(height=0.46/2,width=0.46/2).set_stroke(width=0.).set_fill(current_color, opacity=0.7)
dataset_target.move_to(data)
dataset_target.generate_target()
aligned_edge = ORIGIN
if i % 2 == 0:
old_target = dataset_target.target
buff -= .25
aligned_edge = LEFT
dataset_target.target.next_to(
sampler, buff=buff, direction=UP,
aligned_edge=LEFT
)
else:
dataset_target.target.next_to(
old_target, direction=RIGHT, buff=0.01,
)
new_datasets.append(dataset_target)
first_animations.append(data.animate(run_time=0.5).set_stroke(current_color))
second_animations.append(MoveToTarget(dataset_target, run_time=1.5))
self.play(*first_animations)
self.play(*second_animations)
self.wait()
move_animation = []
for j,gpu in enumerate(gpus):
buff = 0
for i,data in enumerate(new_datasets):
if i % 2 == 0:
current_color = colors[i//2]
if j != 3:
data = data.copy()
data.generate_target()
aligned_edge = ORIGIN
if i % 2 == 0:
old_target = data.target
buff -= .25
aligned_edge = LEFT
data.target.next_to(
gpu, buff=buff, direction=UP,
aligned_edge=LEFT
)
else:
data.target.next_to(
old_target, direction=RIGHT, buff=0.01,
)
move_animation.append(MoveToTarget(data, run_time=1.5))
self.play(*move_animation)
self.remove(step_1)
step_2 = MarkupText(
f"This behavior is undesireable, because we want\neach GPU to see different data for efficient training.",
font_size=18
)
step_2.move_to([0, -2.5, 0])
self.play(
Write(step_2, run_time=2.5),
)
self.wait()

34
manim_animations/dataloaders/stage_3.py Normal file
View File

@ -0,0 +1,34 @@
# Copyright 2024 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from manim import *
class Stage3(Scene):
def construct(self):
step_1 = MarkupText(
f"To combat this, Accelerate employs one of two different\nSampler wrapper methods depending on the scenario:",
font_size=24
)
step_1.move_to([0, 1.5, 0])
self.add(step_1)
step_2 = MarkupText(
f"1. Sharding the dataset before drawing:\n\t● <span fgcolor='{RED}'>IterableDatasetShard</span>\n\t● <span fgcolor='{RED}'>BatchSamplerShard</span>",
font_size=24,
).next_to(step_1, direction=DOWN, aligned_edge=LEFT)
self.add(step_2)
step_3 = MarkupText(
f"\n\n2. Splitting the batch after drawing:\n\t● <span fgcolor='{BLUE}'>DataLoaderDispatcher</span>",
font_size=24,
).next_to(step_2, direction=DOWN, aligned_edge=LEFT)
self.add(step_3)

52
manim_animations/dataloaders/stage_4.py Normal file
View File

@ -0,0 +1,52 @@
# Copyright 2024 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from manim import *
class Stage4(Scene):
def construct(self):
step_1 = MarkupText(
f"To understand the next part fully, let's define two terms,\n<span fgcolor='{RED}'>`batch_size`</span> and <span fgcolor='{BLUE}'>`global_batch_size`</span>:",
font_size=18
)
step_1.move_to([0, 1.5, 0])
# <span fgcolor='{YELLOW}'>●</span>
step_2 = MarkupText(
f"\n\n● <span fgcolor='{RED}'>`batch_size`</span>: \n\tThis will be defined as the batch size seen on a given\n\t*individual* GPU",
font_size=18,
).next_to(step_1, direction=DOWN, aligned_edge=LEFT)
step_3 = MarkupText(
f"\n\n● <span fgcolor='{BLUE}'>`global_batch_size`</span>:\n\tThis will be defined as the *total* number of\n\tdifferent items seen in the dataset, across all GPUs",
font_size=18,
).next_to(step_2, direction=DOWN, aligned_edge=LEFT)
step_4 = MarkupText(
f"\n\nSo if we have a dataset of 64 items, 8 GPUs, \nand a `batch_size` of 8, each *step* will go through\nthe entire dataset one time as 8*8=64",
font_size=18,
).next_to(step_3, direction=DOWN, aligned_edge=LEFT)
self.play(
Write(step_1, run_time=4),
)
self.play(
Write(step_2, run_time=4)
)
self.play(
Write(step_3, run_time=4)
)
self.play(
Write(step_4, run_time=6)
)
self.wait()

203
manim_animations/dataloaders/stage_5.py Normal file
View File

@ -0,0 +1,203 @@
# Copyright 2024 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from manim import *
class Stage5(Scene):
def construct(self):
# The dataset items
colors = ["BLUE_E", "DARK_BROWN", "GOLD_E", "GRAY_A"]
fill = Rectangle(height=0.46,width=0.46).set_stroke(width=0)
columns = [
VGroup(*[Rectangle(height=0.25,width=0.25,color=colors[j]) for i in range(8)]).arrange(RIGHT,buff=0)
for j in range(4)
]
dataset_recs = VGroup(*columns).arrange(UP, buff=0)
dataset_text = Text("Dataset", font_size=24)
dataset = Group(dataset_recs,dataset_text).arrange(DOWN, buff=0.5, aligned_edge=DOWN)
dataset.move_to([-2,0,0])
self.add(dataset)
code = Code(
code="# We enable this by default\naccelerator = Accelerator()\ndataloader = DataLoader(...)\ndataloader = accelerator.prepare(dataloader)\nfor batch in dataloader:\n\t...",
tab_width=4,
background="window",
language="Python",
font="Monospace",
font_size=14,
corner_radius=.2,
insert_line_no=False,
line_spacing=.75,
style=Code.styles_list[1],
)
code.move_to([-3.5, 2.5, 0])
self.add(code)
# The dataloader itself
sampler_1 = Group(
Rectangle(color="blue", height=1, width=1),
Text("Sampler GPU 1", font_size=12)
).arrange(DOWN, buff=.25, aligned_edge=DOWN)
sampler_2 = Group(
Rectangle(color="blue", height=1, width=1),
Text("Sampler GPU 2", font_size=12)
).arrange(DOWN, buff=.25, aligned_edge=DOWN)
sampler_3 = Group(
Rectangle(color="blue", height=1, width=1),
Text("Sampler GPU 3", font_size=12)
).arrange(DOWN, buff=.25, aligned_edge=DOWN)
sampler_4 = Group(
Rectangle(color="blue", height=1, width=1),
Text("Sampler GPU 4", font_size=12)
).arrange(DOWN, buff=.25, aligned_edge=DOWN)
sampler_1.move_to([2,2,0])
sampler_2.move_to([2,.5,0])
sampler_3.move_to([2,-1.,0])
sampler_4.move_to([2,-2.5,0])
self.add(sampler_1, sampler_2, sampler_3, sampler_4)
samplers = [sampler_1[0], sampler_2[0], sampler_3[0], sampler_4[0]]
gpu_1 = Group(
Rectangle(color="white", height=1, width=1),
Text("Output GPU 1", font_size=12)
).arrange(DOWN, buff=.25, aligned_edge=DOWN).move_to([4.5, 2, 0])
gpu_2 = Group(
Rectangle(color="white", height=1, width=1),
Text("Output GPU 2", font_size=12)
).arrange(DOWN, buff=.25, aligned_edge=DOWN).move_to([4.5, .5, 0])
gpu_3 = Group(
Rectangle(color="white", height=1, width=1),
Text("Output GPU 3", font_size=12)
).arrange(DOWN, buff=.25, aligned_edge=DOWN).move_to([4.5, -1, 0])
gpu_4 = Group(
Rectangle(color="white", height=1, width=1),
Text("Output GPU 4", font_size=12)
).arrange(DOWN, buff=.25, aligned_edge=DOWN).move_to([4.5, -2.5, 0])
gpus = [gpu_1[0], gpu_2[0], gpu_3[0], gpu_4[0]]
self.add(gpu_1, gpu_2, gpu_3, gpu_4)
# Animate their existence
self.play(
Create(gpu_1[0], run_time=1),
Create(gpu_2[0], run_time=1),
Create(gpu_3[0], run_time=1),
Create(gpu_4[0], run_time=1),
Create(dataset_recs, run_time=1),
Create(sampler_1[0], run_time=1),
Create(sampler_2[0], run_time=1),
Create(sampler_3[0], run_time=1),
Create(sampler_4[0], run_time=1),
)
first_animations = []
second_animations = []
colors = ["BLUE_E", "DARK_BROWN", "GOLD_E", "GRAY_A"]
current_color = colors[0]
buff = 0
lr_buff = .25
old_target = None
new_datasets = []
for i,row_data in enumerate(dataset_recs):
new_row = []
current_color = colors[i]
if i == 0:
idx = -3
elif i == 1:
idx = -2
elif i == 2:
idx = -1
elif i == 3:
idx = 0
for j,indiv_data in enumerate(row_data):
dataset_target = Rectangle(height=0.46/2,width=0.46/2).set_stroke(width=0.).set_fill(current_color, opacity=0.7)
dataset_target.move_to(indiv_data)
dataset_target.generate_target()
aligned_edge = ORIGIN
if j % 8 == 0:
aligned_edge = LEFT
dataset_target.target.next_to(
samplers[abs(idx)].get_corner(UP+LEFT), buff=.02, direction=RIGHT+DOWN,
)
dataset_target.target.set_x(dataset_target.target.get_x())
elif j % 4 == 0:
old_target = dataset_target.target
dataset_target.target.next_to(
samplers[abs(idx)].get_corner(UP+LEFT), buff=.02, direction=RIGHT+DOWN,
)
dataset_target.target.set_x(dataset_target.target.get_x())
dataset_target.target.set_y(dataset_target.target.get_y()-.25)
else:
dataset_target.target.next_to(
old_target, direction=RIGHT, buff=0.02,
)
old_target = dataset_target.target
new_row.append(dataset_target)
first_animations.append(indiv_data.animate(run_time=0.5).set_stroke(current_color))
second_animations.append(MoveToTarget(dataset_target, run_time=1.5))
new_datasets.append(new_row)
step_1 = MarkupText(
f"Since we splice the dataset between each GPU,\nthe models weights can be averaged during `backward()`\nActing as though we did one giant epoch\nvery quickly.",
font_size=18
)
step_1.move_to([-2.5, -2, 0])
self.play(
Write(step_1, run_time=3),
)
self.play(
*first_animations,
)
self.play(*second_animations)
self.wait(duration=.5)
move_animation = []
import random
for i,row in enumerate(new_datasets):
# row = [row[k] for k in random.sample(range(8), 8)]
current_color = colors[i]
if i == 0:
idx = -3
elif i == 1:
idx = -2
elif i == 2:
idx = -1
elif i == 3:
idx = 0
for j,indiv_data in enumerate(row):
indiv_data.generate_target()
aligned_edge = ORIGIN
if j % 8 == 0:
aligned_edge = LEFT
indiv_data.target.next_to(
gpus[abs(idx)].get_corner(UP+LEFT), buff=.02, direction=RIGHT+DOWN,
)
indiv_data.target.set_x(indiv_data.target.get_x())
elif j % 4 == 0:
indiv_data.target.next_to(
gpus[abs(idx)].get_corner(UP+LEFT), buff=.02, direction=RIGHT+DOWN,
)
indiv_data.target.set_x(indiv_data.target.get_x())
indiv_data.target.set_y(indiv_data.target.get_y()-.25)
else:
indiv_data.target.next_to(
old_target, direction=RIGHT, buff=0.02,
)
old_target = indiv_data.target
move_animation.append(MoveToTarget(indiv_data, run_time=1.5))
self.play(*move_animation)
self.wait()

193
manim_animations/dataloaders/stage_6.py Normal file
View File

@ -0,0 +1,193 @@
# Copyright 2024 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from manim import *
class Stage6(Scene):
def construct(self):
# The dataset items
colors = ["BLUE_E", "DARK_BROWN", "GOLD_E", "GRAY_A"]
fill = Rectangle(height=0.46,width=0.46).set_stroke(width=0)
columns = [
VGroup(*[Rectangle(height=0.25,width=0.25,color=colors[j]) for i in range(8)]).arrange(RIGHT,buff=0)
for j in range(4)
]
dataset_recs = VGroup(*columns).arrange(UP, buff=0)
dataset_text = Text("Dataset", font_size=24)
dataset = Group(dataset_recs,dataset_text).arrange(DOWN, buff=0.5, aligned_edge=DOWN)
dataset.move_to([-2,0,0])
self.add(dataset)
code = Code(
code="# We enable this by default\naccelerator = Accelerator()\ndataloader = DataLoader(..., shuffle=True)\ndataloader = accelerator.prepare(dataloader)\nfor batch in dataloader:\n\t...",
tab_width=4,
background="window",
language="Python",
font="Monospace",
font_size=14,
corner_radius=.2,
insert_line_no=False,
line_spacing=.75,
style=Code.styles_list[1],
)
code.move_to([-3.5, 2.5, 0])
self.add(code)
# The dataloader itself
sampler_1 = Group(
Rectangle(color="blue", height=1, width=1),
Text("Sampler GPU 1", font_size=12)
).arrange(DOWN, buff=.25, aligned_edge=DOWN)
sampler_2 = Group(
Rectangle(color="blue", height=1, width=1),
Text("Sampler GPU 2", font_size=12)
).arrange(DOWN, buff=.25, aligned_edge=DOWN)
sampler_3 = Group(
Rectangle(color="blue", height=1, width=1),
Text("Sampler GPU 3", font_size=12)
).arrange(DOWN, buff=.25, aligned_edge=DOWN)
sampler_4 = Group(
Rectangle(color="blue", height=1, width=1),
Text("Sampler GPU 4", font_size=12)
).arrange(DOWN, buff=.25, aligned_edge=DOWN)
sampler_1.move_to([2,2,0])
sampler_2.move_to([2,.5,0])
sampler_3.move_to([2,-1.,0])
sampler_4.move_to([2,-2.5,0])
self.add(sampler_1, sampler_2, sampler_3, sampler_4)
samplers = [sampler_1[0], sampler_2[0], sampler_3[0], sampler_4[0]]
gpu_1 = Group(
Rectangle(color="white", height=1, width=1),
Text("Output GPU 1", font_size=12)
).arrange(DOWN, buff=.25, aligned_edge=DOWN).move_to([4.5, 2, 0])
gpu_2 = Group(
Rectangle(color="white", height=1, width=1),
Text("Output GPU 2", font_size=12)
).arrange(DOWN, buff=.25, aligned_edge=DOWN).move_to([4.5, .5, 0])
gpu_3 = Group(
Rectangle(color="white", height=1, width=1),
Text("Output GPU 3", font_size=12)
).arrange(DOWN, buff=.25, aligned_edge=DOWN).move_to([4.5, -1, 0])
gpu_4 = Group(
Rectangle(color="white", height=1, width=1),
Text("Output GPU 4", font_size=12)
).arrange(DOWN, buff=.25, aligned_edge=DOWN).move_to([4.5, -2.5, 0])
gpus = [gpu_1[0], gpu_2[0], gpu_3[0], gpu_4[0]]
self.add(gpu_1, gpu_2, gpu_3, gpu_4)
first_animations = []
second_animations = []
colors = ["BLUE_E", "DARK_BROWN", "GOLD_E", "GRAY_A"]
current_color = colors[0]
buff = 0
lr_buff = .25
old_target = None
new_datasets = []
for i,row_data in enumerate(dataset_recs):
new_row = []
current_color = colors[i]
if i == 0:
idx = -3
elif i == 1:
idx = -2
elif i == 2:
idx = -1
elif i == 3:
idx = 0
for j,indiv_data in enumerate(row_data):
dataset_target = Rectangle(height=0.46/2,width=0.46/2).set_stroke(width=0.).set_fill(current_color, opacity=0.7)
dataset_target.move_to(indiv_data)
dataset_target.generate_target()
aligned_edge = ORIGIN
if j % 8 == 0:
aligned_edge = LEFT
old_target = dataset_target.target
dataset_target.target.next_to(
samplers[abs(idx)].get_corner(UP+LEFT), buff=.02, direction=RIGHT+DOWN,
)
dataset_target.target.set_x(dataset_target.target.get_x())
elif j % 4 == 0:
old_target = dataset_target.target
dataset_target.target.next_to(
samplers[abs(idx)].get_corner(UP+LEFT), buff=.02, direction=RIGHT+DOWN,
)
dataset_target.target.set_x(dataset_target.target.get_x())
dataset_target.target.set_y(dataset_target.target.get_y()-.25)
else:
dataset_target.target.next_to(
old_target, direction=RIGHT, buff=0.02,
)
old_target = dataset_target.target
new_row.append(dataset_target)
first_animations.append(indiv_data.animate(run_time=0.5).set_stroke(current_color))
second_animations.append(MoveToTarget(dataset_target, run_time=1.5))
new_datasets.append(new_row)
step_1 = MarkupText(
f"During shuffling, each mini-batch's\noutput order will be modified",
font_size=18
)
step_1.move_to([-1.5, -2, 0])
self.play(
Write(step_1, run_time=3),
)
self.play(
*first_animations,
)
self.play(*second_animations)
self.wait(duration=.5)
move_animation = []
import random
for i,row in enumerate(new_datasets):
row = [row[k] for k in random.sample(range(8), 8)]
current_color = colors[i]
if i == 0:
idx = -3
elif i == 1:
idx = -2
elif i == 2:
idx = -1
elif i == 3:
idx = 0
for j,indiv_data in enumerate(row):
indiv_data.generate_target()
aligned_edge = ORIGIN
if j % 8 == 0:
aligned_edge = LEFT
indiv_data.target.next_to(
gpus[abs(idx)].get_corner(UP+LEFT), buff=.02, direction=RIGHT+DOWN,
)
indiv_data.target.set_x(indiv_data.target.get_x())
elif j % 4 == 0:
indiv_data.target.next_to(
gpus[abs(idx)].get_corner(UP+LEFT), buff=.02, direction=RIGHT+DOWN,
)
indiv_data.target.set_x(indiv_data.target.get_x())
indiv_data.target.set_y(indiv_data.target.get_y()-.25)
else:
indiv_data.target.next_to(
old_target, direction=RIGHT, buff=0.02,
)
old_target = indiv_data.target
move_animation.append(MoveToTarget(indiv_data, run_time=1.5))
self.play(*move_animation)
self.wait()

182
manim_animations/dataloaders/stage_7.py Normal file
View File

@ -0,0 +1,182 @@
# Copyright 2024 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from manim import *
class Stage7(Scene):
def construct(self):
# The dataset items
code = Code(
code="accelerator = Accelerator(dispatch_batches=True)\ndataloader = DataLoader(...)\ndataloader = accelerator.prepare(dataloader)\nfor batch in dataloader:\n\t...",
tab_width=4,
background="window",
language="Python",
font="Monospace",
font_size=14,
corner_radius=.2,
insert_line_no=False,
line_spacing=.75,
style=Code.styles_list[1],
)
code.move_to([-3.5, 2.5, 0])
self.add(code)
colors = ["BLUE_E", "DARK_BROWN", "GOLD_E", "GRAY_A"]
fill = Rectangle(height=0.46,width=0.46).set_stroke(width=0)
columns = [
VGroup(*[Rectangle(height=0.25,width=0.25,color=colors[j]) for i in range(8)]).arrange(RIGHT,buff=0)
for j in range(4)
]
dataset_recs = VGroup(*columns).arrange(UP, buff=0)
dataset_text = Text("Dataset", font_size=24)
dataset = Group(dataset_recs,dataset_text).arrange(DOWN, buff=0.5, aligned_edge=DOWN)
dataset.move_to([-2,0,0])
self.add(dataset)
# The dataloader itself
sampler_1 = Group(
Rectangle(color="blue", height=1.02, width=1.02),
Text("Sampler GPU 1", font_size=12)
).arrange(DOWN, buff=.25, aligned_edge=DOWN)
sampler_2 = Group(
Rectangle(color="blue", height=1.02, width=1.02),
Text("Sampler GPU 2", font_size=12)
).arrange(DOWN, buff=.25, aligned_edge=DOWN)
sampler_3 = Group(
Rectangle(color="blue", height=1.02, width=1.02),
Text("Sampler GPU 3", font_size=12)
).arrange(DOWN, buff=.25, aligned_edge=DOWN)
sampler_4 = Group(
Rectangle(color="blue", height=1.02, width=1.02),
Text("Sampler GPU 4", font_size=12)
).arrange(DOWN, buff=.25, aligned_edge=DOWN)
sampler_1.move_to([2,2,0])
sampler_2.move_to([2,.5,0])
sampler_3.move_to([2,-1.,0])
sampler_4.move_to([2,-2.5,0])
self.add(sampler_1, sampler_2, sampler_3, sampler_4)
samplers = [sampler_1[0], sampler_2[0], sampler_3[0], sampler_4[0]]
gpu_1 = Group(
Rectangle(color="white", height=1.02, width=.98),
Text("Output GPU 1", font_size=12)
).arrange(DOWN, buff=.25, aligned_edge=DOWN).move_to([4.5, 2, 0])
gpu_2 = Group(
Rectangle(color="white", height=1.02, width=.98),
Text("Output GPU 2", font_size=12)
).arrange(DOWN, buff=.25, aligned_edge=DOWN).move_to([4.5, .5, 0])
gpu_3 = Group(
Rectangle(color="white", height=1.02, width=.98),
Text("Output GPU 3", font_size=12)
).arrange(DOWN, buff=.25, aligned_edge=DOWN).move_to([4.5, -1, 0])
gpu_4 = Group(
Rectangle(color="white", height=1.02, width=.98),
Text("Output GPU 4", font_size=12)
).arrange(DOWN, buff=.25, aligned_edge=DOWN).move_to([4.5, -2.5, 0])
gpus = [gpu_1[0], gpu_2[0], gpu_3[0], gpu_4[0]]
self.add(gpu_1, gpu_2, gpu_3, gpu_4)
step_1 = MarkupText(
f"When using a `DataLoaderDispatcher`, all\nof the samples are collected from GPU 0's dataset,\nthen divided and sent to each GPU.\nAs a result, this will be slower.",
font_size=18
)
step_1.move_to([-2.5, -2, 0])
self.play(
Write(step_1, run_time=3.5),
)
first_animations = []
second_animations = []
colors = ["BLUE_E", "DARK_BROWN", "GOLD_E", "GRAY_A"]
current_color = colors[0]
ud_buff = 0.01
lr_buff = 0.01
old_target = None
new_datasets = []
for i,row_data in enumerate(dataset_recs):
new_row = []
current_color = colors[i]
for j,indiv_data in enumerate(row_data):
dataset_target = Rectangle(height=0.46/4,width=0.46/2).set_stroke(width=0.).set_fill(current_color, opacity=0.7)
dataset_target.move_to(indiv_data)
dataset_target.generate_target()
aligned_edge = ORIGIN
if j % 8 == 0:
aligned_edge = LEFT
dataset_target.target.next_to(
samplers[0].get_corner(DOWN+LEFT), buff=0.0125, direction=RIGHT+UP,
)
dataset_target.target.set_x(dataset_target.target.get_x())
dataset_target.target.set_y(dataset_target.target.get_y() + (.25 * i))
elif j % 4 == 0:
old_target = dataset_target.target
dataset_target.target.next_to(
samplers[0].get_corner(DOWN+LEFT), buff=0.0125, direction=RIGHT+UP,
)
dataset_target.target.set_x(dataset_target.target.get_x())
dataset_target.target.set_y(dataset_target.target.get_y()+.125 + (.25 * i))
else:
dataset_target.target.next_to(
old_target, direction=RIGHT, buff=0.0125,
)
old_target = dataset_target.target
new_row.append(dataset_target)
first_animations.append(indiv_data.animate(run_time=0.5).set_stroke(current_color))
second_animations.append(MoveToTarget(dataset_target, run_time=1.5))
new_datasets.append(new_row)
self.play(
*first_animations,
)
self.play(*second_animations)
move_animation = []
for i,row in enumerate(new_datasets):
current_color = colors[i]
if i == 0:
idx = -3
elif i == 1:
idx = -2
elif i == 2:
idx = -1
elif i == 3:
idx = 0
for j,indiv_data in enumerate(row):
indiv_data.generate_target()
indiv_data.animate.stretch_to_fit_height(0.46/2)
aligned_edge = ORIGIN
if j % 8 == 0:
aligned_edge = LEFT
indiv_data.target.next_to(
gpus[abs(idx)].get_corner(UP+LEFT), buff=.01, direction=RIGHT+DOWN,
)
indiv_data.target.set_x(indiv_data.target.get_x())
indiv_data.target.set_y(indiv_data.target.get_y()-.25)
elif j % 4 == 0:
indiv_data.target.next_to(
gpus[abs(idx)].get_corner(UP+LEFT), buff=.01, direction=RIGHT+DOWN,
)
indiv_data.target.set_x(indiv_data.target.get_x())
else:
indiv_data.target.next_to(
old_target, direction=RIGHT, buff=0.01,
)
old_target = indiv_data.target
move_animation.append(MoveToTarget(indiv_data, run_time=1.5))
self.play(*move_animation)
self.wait()