opensource/ManimML
1
0
Fork 0
mirror of https://github.com/helblazer811/ManimML.git synced 2025-06-19 20:59:03 +08:00
Code Issues Packages Projects Releases Wiki Activity

Reformatted the code using black, allowd for different orientation NNs, made an option for highlighting the active filter in a CNN forward pass.

Browse Source
This commit is contained in:
Alec Helbling
2023-01-09 15:52:37 +09:00
parent 39b0b133ce
commit ba63116b37
19 changed files with 485 additions and 283 deletions
Download Patch File Download Diff File Expand all files Collapse all files

111
manim_ml/neural_network/animations/dropout.py
View File

@ -6,10 +6,12 @@ from manim import *
import random
from manim_ml.neural_network.layers.feed_forward import FeedForwardLayer
from manim_ml.neural_network.layers.feed_forward_to_feed_forward import FeedForwardToFeedForward
from manim_ml.neural_network.layers.feed_forward_to_feed_forward import (
FeedForwardToFeedForward,
)
class XMark(VGroup):
def __init__(self, stroke_width=1.0, color=GRAY):
super().__init__()
line_one = Line(
@ -17,7 +19,7 @@ class XMark(VGroup):
[0.1, -0.1, 0],
stroke_width=1.0,
stroke_color=color,
z_index=4
z_index=4,
)
self.add(line_one)
line_two = Line(
@ -25,10 +27,11 @@ class XMark(VGroup):
[0.1, 0.1, 0],
stroke_width=1.0,
stroke_color=color,
z_index=4
z_index=4,
)
self.add(line_two)
def get_edges_to_drop_out(layer: FeedForwardToFeedForward, layers_to_nodes_to_drop_out):
"""Returns edges to drop out for a given FeedForwardToFeedForward layer"""
prev_layer = layer.input_layer
@ -43,18 +46,21 @@ def get_edges_to_drop_out(layer: FeedForwardToFeedForward, layers_to_nodes_to_dr
prev_node_index = int(edge_index / next_layer.num_nodes)
next_node_index = edge_index % next_layer.num_nodes
# Check if the edges should be dropped out
if prev_node_index in prev_layer_nodes_to_dropout \
or next_node_index in next_layer_nodes_to_dropout:
if (
prev_node_index in prev_layer_nodes_to_dropout
or next_node_index in next_layer_nodes_to_dropout
):
edges_to_dropout.append(edge)
edge_indices_to_dropout.append(edge_index)
return edges_to_dropout, edge_indices_to_dropout
def make_pre_dropout_animation(
neural_network,
layers_to_nodes_to_drop_out,
dropped_out_color=GRAY,
dropped_out_opacity=0.2
dropped_out_opacity=0.2,
):
"""Makes an animation that sets up the NN layer for dropout"""
animations = []
@ -70,8 +76,7 @@ def make_pre_dropout_animation(
layers_to_edges_to_dropout = {}
for layer in feed_forward_to_feed_forward_layers:
layers_to_edges_to_dropout[layer], _ = get_edges_to_drop_out(
layer,
layers_to_nodes_to_drop_out
layer, layers_to_nodes_to_drop_out
)
# Dim the colors of the edges
dim_edge_colors_animations = []
@ -92,12 +97,9 @@ def make_pre_dropout_animation(
)
"""
dim_edge_colors_animations.append(
FadeOut(edge)
)
dim_edge_colors_animations.append(FadeOut(edge))
dim_edge_colors_animation = AnimationGroup(
*dim_edge_colors_animations,
lag_ratio=0.0
*dim_edge_colors_animations, lag_ratio=0.0
)
# Dim the colors of the nodes
dim_nodes_animations = []
@ -113,10 +115,7 @@ def make_pre_dropout_animation(
create_x = Create(x_mark)
dim_nodes_animations.append(create_x)
dim_nodes_animation = AnimationGroup(
*dim_nodes_animations,
lag_ratio=0.0
)
dim_nodes_animation = AnimationGroup(*dim_nodes_animations, lag_ratio=0.0)
animation_group = AnimationGroup(
dim_edge_colors_animation,
@ -125,6 +124,7 @@ def make_pre_dropout_animation(
return animation_group, x_marks
def make_post_dropout_animation(
neural_network,
layers_to_nodes_to_drop_out,
@ -143,19 +143,15 @@ def make_post_dropout_animation(
layers_to_edges_to_dropout = {}
for layer in feed_forward_to_feed_forward_layers:
layers_to_edges_to_dropout[layer], _ = get_edges_to_drop_out(
layer,
layers_to_nodes_to_drop_out
layer, layers_to_nodes_to_drop_out
)
# Remove the x marks
# Remove the x marks
uncreate_animations = []
for x_mark in x_marks:
uncreate_x_mark = Uncreate(x_mark)
uncreate_animations.append(uncreate_x_mark)
uncreate_x_marks = AnimationGroup(
*uncreate_animations,
lag_ratio=0.0
)
uncreate_x_marks = AnimationGroup(*uncreate_animations, lag_ratio=0.0)
# Add the edges back
create_edge_animations = []
for layer in layers_to_edges_to_dropout.keys():
@ -164,20 +160,12 @@ def make_post_dropout_animation(
for edge_index, edge in enumerate(edges_to_drop_out):
edge_copy = edge.copy()
edges_to_drop_out[edge_index] = edge_copy
create_edge_animations.append(
FadeIn(edge_copy)
)
create_edge_animations.append(FadeIn(edge_copy))
create_edge_animation = AnimationGroup(
*create_edge_animations,
lag_ratio=0.0
)
create_edge_animation = AnimationGroup(*create_edge_animations, lag_ratio=0.0)
return AnimationGroup(uncreate_x_marks, create_edge_animation, lag_ratio=0.0)
return AnimationGroup(
uncreate_x_marks,
create_edge_animation,
lag_ratio=0.0
)
def make_forward_pass_with_dropout_animation(
neural_network,
@ -195,35 +183,25 @@ def make_forward_pass_with_dropout_animation(
)
# Iterate through network and get feed forward layers
for layer in feed_forward_layers:
layer_args[layer] = {
"dropout_node_indices": layers_to_nodes_to_drop_out[layer]
}
layer_args[layer] = {"dropout_node_indices": layers_to_nodes_to_drop_out[layer]}
for layer in feed_forward_to_feed_forward_layers:
_, edge_indices = get_edges_to_drop_out(
layer,
layers_to_nodes_to_drop_out
)
layer_args[layer] = {
"edge_indices_to_dropout": edge_indices
}
_, edge_indices = get_edges_to_drop_out(layer, layers_to_nodes_to_drop_out)
layer_args[layer] = {"edge_indices_to_dropout": edge_indices}
return neural_network.make_forward_pass_animation(layer_args=layer_args)
return neural_network.make_forward_pass_animation(
layer_args=layer_args
)
def make_neural_network_dropout_animation(
neural_network,
dropout_rate=0.5,
do_forward_pass=True
neural_network, dropout_rate=0.5, do_forward_pass=True
):
"""
Makes a dropout animation for a given neural network.
Makes a dropout animation for a given neural network.
NOTE Does dropout only on feed forward layers.
1. Does dropout
2. If `do_forward_pass` then do forward pass animation
3. Revert network to pre-dropout appearance
NOTE Does dropout only on feed forward layers.
1. Does dropout
2. If `do_forward_pass` then do forward pass animation
3. Revert network to pre-dropout appearance
"""
# Go through the network and get the FeedForwardLayer instances
feed_forward_layers = neural_network.filter_layers(
@ -247,26 +225,19 @@ def make_neural_network_dropout_animation(
layers_to_nodes_to_drop_out[feed_forward_layer] = nodes_to_drop_out
# Make the animation
pre_dropout_animation, x_marks = make_pre_dropout_animation(
neural_network,
layers_to_nodes_to_drop_out
neural_network, layers_to_nodes_to_drop_out
)
if do_forward_pass:
forward_pass_animation = make_forward_pass_with_dropout_animation(
neural_network,
layers_to_nodes_to_drop_out
neural_network, layers_to_nodes_to_drop_out
)
else:
forward_pass_animation = AnimationGroup()
post_dropout_animation = make_post_dropout_animation(
neural_network,
layers_to_nodes_to_drop_out,
x_marks
neural_network, layers_to_nodes_to_drop_out, x_marks
)
# Combine the animations into one
return Succession(
pre_dropout_animation,
forward_pass_animation,
post_dropout_animation
pre_dropout_animation, forward_pass_animation, post_dropout_animation
)