JupyterLab-compatible for SO https://stackoverflow.com/q/75017358/8508004¶
This builds on the Jupyter notebook with related code that can be viewed here.
Click here to launch that notebook in an active, temporary Jupyter session served via MyBinder.org with the environment already having Python and ALMOST everything needed to make the animation with a widget work in the notebook installed and working. Here, I'm going to include the %pip install torchflow line as a cell so that the notebook can just be run. (If you use [that link above]((https://mybinder.org/v2/gh/fomightez/animated_matplotlib-binder/master?urlpath=git-pull%3Frepo%3Dhttps%253A%252F%252Fgist.github.com%252Ffomightez%252Ff539a3770f2f3287bf12de2d2a549e3a%26urlpath%3Dtree%252Ff539a3770f2f3287bf12de2d2a549e3a%252Fsegments_for_training_updated_animation.ipynb), you'll need to run the same process first.)
Go ahead and run that below to get the final preparation of the environment out of the way.
%pip install torch
Collecting torch
Downloading torch-1.13.1-cp37-cp37m-manylinux1_x86_64.whl (887.5 MB)
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 887.5/887.5 MB 1.2 MB/s eta 0:00:0000:0100:01
Requirement already satisfied: typing-extensions in /srv/conda/envs/notebook/lib/python3.7/site-packages (from torch) (4.4.0)
Collecting nvidia-cuda-nvrtc-cu11==11.7.99
Downloading nvidia_cuda_nvrtc_cu11-11.7.99-2-py3-none-manylinux1_x86_64.whl (21.0 MB)
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 21.0/21.0 MB 53.1 MB/s eta 0:00:0000:0100:01
Collecting nvidia-cublas-cu11==11.10.3.66
Downloading nvidia_cublas_cu11-11.10.3.66-py3-none-manylinux1_x86_64.whl (317.1 MB)
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 317.1/317.1 MB 3.0 MB/s eta 0:00:0000:0100:01
Collecting nvidia-cudnn-cu11==8.5.0.96
Downloading nvidia_cudnn_cu11-8.5.0.96-2-py3-none-manylinux1_x86_64.whl (557.1 MB)
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 557.1/557.1 MB 1.8 MB/s eta 0:00:0000:0100:01
Collecting nvidia-cuda-runtime-cu11==11.7.99
Downloading nvidia_cuda_runtime_cu11-11.7.99-py3-none-manylinux1_x86_64.whl (849 kB)
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 849.3/849.3 kB 1.9 MB/s eta 0:00:0000:0100:01
Requirement already satisfied: setuptools in /srv/conda/envs/notebook/lib/python3.7/site-packages (from nvidia-cublas-cu11==11.10.3.66->torch) (65.6.3)
Requirement already satisfied: wheel in /srv/conda/envs/notebook/lib/python3.7/site-packages (from nvidia-cublas-cu11==11.10.3.66->torch) (0.38.4)
Installing collected packages: nvidia-cuda-runtime-cu11, nvidia-cuda-nvrtc-cu11, nvidia-cublas-cu11, nvidia-cudnn-cu11, torch
Successfully installed nvidia-cublas-cu11-11.10.3.66 nvidia-cuda-nvrtc-cu11-11.7.99 nvidia-cuda-runtime-cu11-11.7.99 nvidia-cudnn-cu11-8.5.0.96 torch-1.13.1
Note: you may need to restart the kernel to use updated packages.
Restart the kernel after running that cell and then try running the one below. With that preparation complete, continue on to examine and run this method...
Right now this method (based on based on here) is more universal than the one illustrated in the first notebook. This method works in both the traditional Jupyter notebook and JupyterLab as of now (January 2023).
In the previous notebook with the other method (here), the data was collected first and then the collected data used to plot using fig.canvas.draw() to reset the image each round. A pause is built in to make sure the animation steps through adding showing each next segment with time in between instead of just quickly blinking through doing that too fast to really see and just displaying the last step.
Here the data is plotted in the course of each iteration, which corresponds to the original loop in the OP's code. The time it takes to make each round of data actually build in more pause than the other method where the data had already been collected. I've added in additional cushion here though to make it even smoother. This additional cushio can be removed or adjusted as needed here. It would be important to have this pause to control running through displaying the data if the data was already existing or not computational time intensive to produce.
In addition to the extra delay, the live_plot() funnction that handles updating the plot with each round uses IPython's display control to clear the output with each frame that will be shown in the animation. You can see a somewhat simpler implementation in the StackOverflow answer that this approach is based on.
Run this code below to see the animation. Note that it doesn't need setting any %matplotlib magic explicitly. This is as opposed to the use of %matplotlib notebook being necessaryfor related code that can be viewed here to work in the tradtional Jupyter notbeook interface.
# based on https://stackoverflow.com/a/52672859/8508004 to help with https://stackoverflow.com/q/75017358/8508004
from IPython.display import clear_output
from matplotlib import pyplot as plt
import numpy as np
import collections
import time
def live_plot(data_dict, figsize=(12,5), title=''):
clear_output(wait=True)
plt.figure(figsize=figsize)
#plt.plot(data_dict["steps"],data_dict["r"] , 'r-', label = "real")
#plt.plot(data_dict["steps"],data_dict["b"] , 'b-', label = "prediction")
for i,_ in enumerate(data_dict["steps"]):
plt.plot(data_dict["steps"][i], list(data_dict["r"][i]) , 'r-', )
plt.plot(data_dict["steps"][i], list(data_dict["b"][i]) , 'b-', )
plt.title(title)
plt.grid(True)
#plt.legend(loc='center left') # the plot evolves to the right
plt.show()
time.sleep(0.2) # extend delay between adding next frame in animation
data = collections.defaultdict(list)
import torch
from torch import nn
# torch.manual_seed(1) # reproducible
# Hyper Parameters
TIME_STEP = 10 # rnn time step
INPUT_SIZE = 1 # rnn input size
LR = 0.02 # learning rate
# data
steps = np.linspace(0, np.pi*2, 100, dtype=np.float32) # float32 for converting torch FloatTensor
x_np = np.sin(steps)
y_np = np.cos(steps)
class RNN(nn.Module):
def __init__(self):
super(RNN, self).__init__()
self.rnn = nn.RNN(
input_size=INPUT_SIZE,
hidden_size=32, # rnn hidden unit
num_layers=1, # number of rnn layer
batch_first=True, # input & output will has batch size as 1s dimension. e.g. (batch, time_step, input_size)
)
self.out = nn.Linear(32, 1)
def forward(self, x, h_state):
# x (batch, time_step, input_size)
# h_state (n_layers, batch, hidden_size)
# r_out (batch, time_step, hidden_size)
r_out, h_state = self.rnn(x, h_state)
outs = [] # save all predictions
for time_step in range(r_out.size(1)): # calculate output for each time step
outs.append(self.out(r_out[:, time_step, :]))
return torch.stack(outs, dim=1), h_state
# instead, for simplicity, you can replace above codes by follows
# r_out = r_out.view(-1, 32)
# outs = self.out(r_out)
# outs = outs.view(-1, TIME_STEP, 1)
# return outs, h_state
# or even simpler, since nn.Linear can accept inputs of any dimension
# and returns outputs with same dimension except for the last
# outs = self.out(r_out)
# return outs
rnn = RNN()
print(rnn)
optimizer = torch.optim.Adam(rnn.parameters(), lr=LR) # optimize all cnn parameters
loss_func = nn.MSELoss()
h_state = None # for initial hidden state
for step in range(100):
start, end = step * np.pi, (step+1)*np.pi # time range
# use sin predicts cos
steps = np.linspace(start, end, TIME_STEP, dtype=np.float32, endpoint=False) # float32 for converting torch FloatTensor
x_np = np.sin(steps)
y_np = np.cos(steps)
x = torch.from_numpy(x_np[np.newaxis, :, np.newaxis]) # shape (batch, time_step, input_size)
y = torch.from_numpy(y_np[np.newaxis, :, np.newaxis])
prediction, h_state = rnn(x, h_state) # rnn output
# !! next step is important !!
h_state = h_state.data # repack the hidden state, break the connection from last iteration
loss = loss_func(prediction, y) # calculate loss
optimizer.zero_grad() # clear gradients for this training step
loss.backward() # backpropagation, compute gradients
optimizer.step() # apply gradients
# plotting
data['steps'].append(list(steps))
data['r'].append(y_np.flatten())
data['b'].append(prediction.data.numpy().flatten())
live_plot(data);
Will the method with FuncAnimation() with associated widget controller also work well in JupyterLab?¶
Now that I pulled the original code apart enough to realize each is a segment and implemented a couple approaches, I also wondered if could use the method with FuncAnimation() with associated widget controller that is illustrated at the bottom of here (and that I had used to answer here, recently) to make something that also would work in JupyterLab. (Note it is not all animations methods involving FuncAnimation() that work in JupyterLab, see here for one that only works in classic notebook at this time. The widget is key.)
# If upon first running, it shows a non-interactive, single static shot of the plot below the interactive one with the widget controller,
# JUST RE-RUN TWICE. Re-run usually fixes that display quirk.
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import animation
from IPython.display import HTML, display
plt.rcParams["animation.html"] = "jshtml"
#plt.ioff() #needed so the second time you run it you get only single plot
import collections
import time
fig = plt.figure(figsize=(12,5))
ax = plt.axes(xlim=(0, 320), ylim=(-1.75, 1.75))
lineplot, = ax.plot([], [], "r-")
lineplot2, = ax.plot([], [], "b-")
import torch
from torch import nn
# torch.manual_seed(1) # reproducible
# Hyper Parameters
TIME_STEP = 10 # rnn time step
INPUT_SIZE = 1 # rnn input size
LR = 0.02 # learning rate
# data
steps = np.linspace(0, np.pi*2, 100, dtype=np.float32) # float32 for converting torch FloatTensor
x_np = np.sin(steps)
y_np = np.cos(steps)
class RNN(nn.Module):
def __init__(self):
super(RNN, self).__init__()
self.rnn = nn.RNN(
input_size=INPUT_SIZE,
hidden_size=32, # rnn hidden unit
num_layers=1, # number of rnn layer
batch_first=True, # input & output will has batch size as 1s dimension. e.g. (batch, time_step, input_size)
)
self.out = nn.Linear(32, 1)
def forward(self, x, h_state):
# x (batch, time_step, input_size)
# h_state (n_layers, batch, hidden_size)
# r_out (batch, time_step, hidden_size)
r_out, h_state = self.rnn(x, h_state)
outs = [] # save all predictions
for time_step in range(r_out.size(1)): # calculate output for each time step
outs.append(self.out(r_out[:, time_step, :]))
return torch.stack(outs, dim=1), h_state
# instead, for simplicity, you can replace above codes by follows
# r_out = r_out.view(-1, 32)
# outs = self.out(r_out)
# outs = outs.view(-1, TIME_STEP, 1)
# return outs, h_state
# or even simpler, since nn.Linear can accept inputs of any dimension
# and returns outputs with same dimension except for the last
# outs = self.out(r_out)
# return outs
rnn = RNN()
print(rnn)
optimizer = torch.optim.Adam(rnn.parameters(), lr=LR) # optimize all cnn parameters
loss_func = nn.MSELoss()
h_state = None # for initial hidden state
data = collections.defaultdict(list)
def init():
global h_state, data
lineplot.set_data([], [])
lineplot2.set_data([], [])
data = collections.defaultdict(list)
return lineplot, #return [lineplot] also works like in https://nbviewer.org/github/raphaelquast/jupyter_notebook_intro/blob/master/jupyter_nb_introduction.ipynb#pre-render-animations-and-export-to-HTML
def animate(i):
global h_state, data
step = i
start, end = step * np.pi, (step+1)*np.pi # time range
# use sin predicts cos
steps = np.linspace(start, end, TIME_STEP, dtype=np.float32, endpoint=False) # float32 for converting torch FloatTensor
x_np = np.sin(steps)
y_np = np.cos(steps)
x = torch.from_numpy(x_np[np.newaxis, :, np.newaxis]) # shape (batch, time_step, input_size)
y = torch.from_numpy(y_np[np.newaxis, :, np.newaxis])
prediction, h_state = rnn(x, h_state) # rnn output
# !! next step is important !!
h_state = h_state.data # repack the hidden state, break the connection from last iteration
loss = loss_func(prediction, y) # calculate loss
optimizer.zero_grad() # clear gradients for this training step
loss.backward() # backpropagation, compute gradients
optimizer.step() # apply gradients
# plotting
data['steps'].append(list(steps))
data['r'].append(y_np.flatten())
data['b'].append(prediction.data.numpy().flatten())
#lineplot.set_data([x], [y])
#lineplot2.set_data([x], [z])
lineplot.set_data(data["steps"],data["r"])
lineplot2.set_data(data["steps"],data["b"])
'''
for i,_ in enumerate(data_dict["steps"]):
plt.plot(data_dict["steps"][i], list(data_dict["r"][i]) , 'r-', )
plt.plot(data_dict["steps"][i], list(data_dict["b"][i]) , 'b-', )
'''
return [lineplot]
anim = animation.FuncAnimation(fig, animate, init_func=init,
frames=100, interval=20, blit=True)
anim
RNN( (rnn): RNN(1, 32, batch_first=True) (out): Linear(in_features=32, out_features=1, bias=True) )
Manually scrubbing back and forth with the slider allows you to choose a point in the building of the plot.
Note, that for some reason I've seen it break the widget normal looping ability at this time. I'm not sure what I did to break it. Manually scrubbing back and forth with the slider did still work even when that happend. I don't know what I broke to make the one above not loop? I tried adding to init() and that didn't seem to help. In fact when it broke it, it would also break it for the one below that is simpler below. Weird
I had seen the glitch I saw was not simply due to including multiple lines because this related, simple code one works to keep looping IN A SEPARATE, or new, NOTEBOOK: (It may or may not work here after running the one above.)
# If upon first running, it shows a non-interactive, single static shot of the plot below the interactive one with the widget controller,
# JUST RE-RUN TWICE. Re-run usually fixes that display quirk.
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import animation
from IPython.display import HTML, display
plt.rcParams["animation.html"] = "jshtml"
plt.ioff() #needed so the second time you run it you get only single plot
fig = plt.figure()
ax = plt.axes(xlim=(0, 4), ylim=(-2, 2))
lineplot, = ax.plot([], [], lw=3)
lineplot2, = ax.plot([], [], "r-")
def init():
lineplot.set_data([], [])
return lineplot, #return [lineplot] also works like in https://nbviewer.org/github/raphaelquast/jupyter_notebook_intro/blob/master/jupyter_nb_introduction.ipynb#pre-render-animations-and-export-to-HTML
def animate(i):
x = np.linspace(0, 4, 1000)
y = np.sin(2 * np.pi * (x - 0.01 * i))
z = np.sin(2.2 * np.pi * (x - 0.31 * i))
lineplot.set_data([x], [y])
lineplot2.set_data([x], [z])
return [lineplot,lineplot2]
anim = animation.FuncAnimation(fig, animate, init_func=init,
frames=200, interval=20, blit=True)
anim