Qiskit (56) —量子機械学習

ついでに量子回路によるレイヤでの順伝播と逆伝播でどれくらいの時間を消費するのか見てみたい。textbook のコードに対しては以下のような改造で良いはず・・・だ:

class HybridFunction(Function):
    @staticmethod
    def forward(ctx, input, quantum_circuit, shift, forward_times, backward_times):
        time_start = time.perf_counter()
        ...
        elapsed = time.perf_counter() - time_start
        forward_times.append(elapsed)
        if not hasattr(ctx, 'backward_times'):
            ctx.backward_times = backward_times

        return result

    @staticmethod
    def backward(ctx, grad_output):
        """ Backward pass computation """
        time_start = time.perf_counter()
        ...
        result = torch.tensor([gradients]).float() * grad_output.float(), None, None, None, None
        elapsed = time.perf_counter() - time_start
        ctx.backward_times.append(elapsed)
        return result

class Hybrid(nn.Module):
    def __init__(self, backend, shots, shift):
        ...
        self.forward_times = []
        self.backward_times = []

    def forward(self, input):
        return HybridFunction.apply(input, self.quantum_circuit, self.shift, self.forward_times, self.backward_times)

    def reset(self):
        self.forward_times = []
        self.backward_times = []

for epoch in range(epochs):
    total_loss = []
    for batch_idx, (data, target) in enumerate(train_loader):
        ...
    ...
    print(f'mean forward_times, {np.mean(model.hybrid.forward_times)} ({len(model.hybrid.forward_times)})')
    print(f'mean backward_times, {np.mean(model.hybrid.backward_times)} ({len(model.hybrid.backward_times)})')
    model.hybrid.reset()

すると

Training [5%]	Loss: -0.7993
mean forward_times, 0.00963353352498416 (200)
mean backward_times, 0.01893993059499735 (200)
Training [10%]	Loss: -0.9326
mean forward_times, 0.009449818629918809 (200)
mean backward_times, 0.01929300983500525 (200)
Training [15%]	Loss: -0.9439
mean forward_times, 0.009592005840004277 (200)
mean backward_times, 0.018860545625025225 (200)
Training [20%]	Loss: -0.9363
mean forward_times, 0.009964517429993975 (200)
mean backward_times, 0.019512393174973112 (200)
...

という感じで、1 バッチにつき、順伝播で 10 ms、逆伝播で 20 ms くらいかかっていることが分かる。

書籍付属のコードで強引にバッチサイズを 1 にして計測してみると以下のような感じに:

Training [2.5%]	Loss: 0.2286
mean forward_times, 0.03484408998999242 (200)
mean backward_times, 1.6172968914750345 (200)
Training [5.0%]	Loss: 0.1449
mean forward_times, 0.03464558973000749 (200)
mean backward_times, 1.6136025506350415 (200)

順伝播が 35 ms で逆伝播が 1600 ms のよう。バッチサイズの変更があまり宜しくないとは言え、逆伝播の時間が妙に長い。

DeprecationWarning: The QuantumCircuit.cu3 method is deprecated as of 0.16.0. It will be removed no earlier than 3 months after the release date. You should use the QuantumCircuit.cu method instead, where cu3(ϴ,φ,λ) = cu(ϴ,φ,λ,0).

なども出ているから、全体的に現時点で最新の API で刷新する必要があるのかもしれない。と思ったけど、そこよりは寧ろ HybridFunction.backward の中で

gradient  = torch.sum((ctx.f(data, params) - res) / delta * grad_output)

が 1 回につき 360 ms くらい食う（これを 45 回くらいイテレーションで回す）のが重いらしい。この時間の大部分は ctx.f(data, params) の部分が関与している。また、Hybrid.forward 内の f が 35 ms くらい食っていて、HybridFunction.forward の中の

def f_each(data, params):
    return torch.tensor([f(torch.flatten(d), params) for d in data], dtype=torch.float64)

で何度も呼び出すので、この f_each がかなり重い関数になる様子。よって、最終的には f の中の QMLCircuit.run がボトルネックになっていそうで、実際 perf_counter で計測するとそのようなのでこれ以上どうにもできないことが分かった。あと、この f_each から参照している f が気になってしまうし*1・・・これ以上の追跡は断念する。

*1:実際に f_each を参照する段階で期待するものがバインドされている保証はどこから来るのだろう・・・