[検証] softmax はどれだけ遅いのか

はじめに

「ゼロから作る Deep Learning」を読んでいると，以下の文が目に入ってきた．

(中略)ソフトマックス関数を適用しても、出力の一番大きいニューロンの場所は変わりません。そのため、ニューラルネットワークが分類を行う際には、出力層のソフトマックス関数を省略することができます。実際の問題では、指数関数の計算は、それなりにコンピュータの計算が必要になるので、出力層のソフトマックス関数は省略するのが一般的です。

斎藤 康毅 著，ゼロから作る Deep Learning，O'Reilly Japan，2016，p70-71 より引用．

実際のところどうなるか検証してみた．

検証

環境

• Windows 11
• Python 3.10.9
• torch 2.1.2
• AMD Ryzen 5 5600X 6-Core Processor 3.70 GHz (念のため)

方法

1000行のベクトルへの演算を1000回計算し，合計時間を計測する．
softmaxだけ，argmaxだけ，softmax+argmax の3通りで検証する．

ソースコード

import torch
from torch import nn
import time


N = 1000 # 特に理由はないが，ImageNet など，1000クラス分類をしている想定
x = torch.randn(N,N,dtype=torch.float64)
m = nn.Softmax(dim=0)

# softmax のみ
t_start = time.time()
for i in range(N):
    y = m(x[i])
t_end = time.time()

print("softmax :",(t_end - t_start)*1000,"[ms]")

# argmax のみ
t_start = time.time()
for i in range(N):
    index_y = torch.argmax(x[i])
t_end = time.time()

print("argmax :",(t_end - t_start)*1000,"[ms]")

# softmax + argmax
t_start = time.time()
for i in range(N):
    y = m(x[i])
    index_y = torch.argmax(y[i])
t_end = time.time()

print("softmax + argmax : ",(t_end - t_start)*1000,"[ms]")

結果

softmax : 14.510393142700195 [ms]
argmax : 11.507749557495117 [ms]
softmax + argmax :  18.867015838623047 [ms]

考察

たしかに微増しているが，そこまで神経質になるほどかは疑問である．
何度かプログラムを実行しても計算時間は安定しなかったので，別の原因がありそう．
おそらくcpuのキャッシュとか？(適当)

参考文献

• https://qiita.com/umapyoi/items/7f0de31c2ac89d3c8338
• https://pytorch.org/docs/stable/generated/torch.nn.Softmax.html
• https://pytorch.org/docs/stable/generated/torch.argmax.html
• https://pytorch.org/docs/stable/generated/torch.randn.html
• Ryzen 5 5600Xの性能・スペック : https://www.dospara.co.jp/5info/cts_str_parts_ryzen5.html
• 斎藤 康毅 著，ゼロから作る Deep Learning，O'Reilly Japan，2016