matplotlibのanimationで一枚ごとに違うタイトルを付けたい

機械学習の勉強をしています
線形回帰が進む様子をmatplotlibのanimationにしてみています．
epoch数をタイトルにいれたいのですが，表示されません．

具体的には以下のコードです．

import torch
import torch.nn as nn
import numpy as np

import matplotlib.pyplot as plt
import matplotlib.animation as animation
from matplotlib.animation import PillowWriter

# hyper parameters
input_size = 1
output_size = 1
num_epochs = 200
learning_rate = 0.001

# toy dataset
# 15 samples, 1 features
x_train = np.array([3.3, 4.4, 5.5, 6.71, 6.93, 4.168, 9.779, 6.182, 7.59, 2.167,
                    7.042, 10.791, 5.313, 7.997, 3.1], dtype=np.float32)
y_train = np.array([1.7, 2.76, 2.09, 3.19, 1.694, 1.573, 3.366, 2.596, 2.53, 1.221,
                    2.827, 3.465, 1.65, 2.904, 1.3], dtype=np.float32)
x_train = x_train.reshape(15,1)
y_train = y_train.reshape(15,1)

# linear regression model
class LinearRegression(nn.Module):

    def __init__(self, input_size, output_size):
        super(LinearRegression, self).__init__()
        self.linear = nn.Linear(input_size, output_size)

    def forward(self, x):
        out = self.linear(x)
        return out

model = LinearRegression(input_size, output_size)

# loss and optimizer
# loss function mean squared error
# optimizer stochastic gradient descent
criterion = nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)


# Figure Setting
fig = plt.figure()
ims = []

# train the model
for epoch in range(num_epochs):
   inputs = torch.from_numpy(x_train)
   targets = torch.from_numpy(y_train)

   optimizer.zero_grad()
   outputs = model(inputs)
   loss = criterion(outputs, targets)
   loss.backward()
   optimizer.step()

   if (epoch + 1) % 10 == 0:
       print('Epoch [%d/%d], Loss: %.4f' % (epoch + 1, num_epochs, loss.item()))
       
       predicted = model(torch.from_numpy(x_train)).detach().numpy()
       line, = plt.plot(x_train,predicted,"skyblue")
       tm = plt.title("Epoch = {0}".format(epoch+1))
       ims.append([line,tm])

# save the model
torch.save(model.state_dict(), "model.pkl")

od, = plt.plot(x_train,y_train,"ro")
#plt.title("test")
plt.legend([od,line],["Original Data","Fitted Line"])

ani = animation.ArtistAnimation(fig,ims,interval=50,blit=True,repeat_delay=1000)
plt.show()

print("Save Animation? [y/n]")
should_save_animation = str(input())

if should_save_animation == "y":
    anim.save("anim.gif",writer="pillow")

ArtistAnimationに渡すArtistのリストに問題があると思うのですが，解決策が分かりません．
よろしくお願いいたします．

行動規範の内容に同意します

回答1件

ベストアンサー

FuncAnimation を使うとタイトルも動的に変更できます。

python
1import matplotlib.pyplot as plt
2import numpy as np
3from matplotlib.animation import FuncAnimation
4
5fig = plt.figure(figsize=(4, 3))
6
7def plot(data):
8    plt.cla()
9    x = np.linspace(0, np.pi * 4, 100)
10    y = np.sin(x + data)
11    plt.plot(x, y, c='b')
12    plt.title('data={}'.format(data))
13
14# アニメーションを作成する。
15anim = FuncAnimation(fig, plot, frames=100)
16anim.save('animation.gif', writer='imagemagick')
17
18# Figure を表示する。
19fig.show()

追記

Axes.set_title() の返り値 matplotlib.text.Text を Artists として入れてもタイトルはアニメーションできないようです。
代わりに Axes.text() の返り値ならできるので、これでタイトルをアニメーションできました。

参考
matplotlib artist animation : title or text not changing

python
1# Notebook で inline 表示する場合は以下が必要
2%matplotlib nbagg
3
4#import packages
5import matplotlib.pyplot as plt
6import numpy as np
7from matplotlib.animation import ArtistAnimation
8
9x = np.linspace(0, np.pi * 4, 100)
10
11fig, ax = plt.subplots(figsize=(4, 4))
12frames = []  # 各フレームを構成する Artist 一覧
13for delta in np.linspace(0, np.pi, 30):
14    y = np.sin(x + delta)
15    
16    # 折れ線グラフを作成する。
17    lines = ax.plot(x, y, c='b')
18    title = ax.text(0.5, 1.01, 'delta={:.2f}'.format(delta),
19                     ha='center', va='bottom',
20                     transform=ax.transAxes, fontsize='large')
21    frames.append(lines + [title])
22
23# アニメーションを作成する。
24anim = ArtistAnimation(fig, frames, interval=500)
25
26# gif 画像として保存する。
27anim.save('animation.gif', writer='imagemagick')
28
29# Figure を表示する。
30fig.show()