DDPG by gymnasium 12日目

GPUをニューラルネットワークのテンソル演算に利用します。

Contents

GPUへの飛ばし方

#変数deviceを’cuda’にする

device = torch.device(“cuda” if torch.cuda.is_available() else “cpu”)

#ネットワークのインスタンスを.to(‘cuda’)する

net_gpu.to(device)

# ネットワークへの入力xを.to(‘cuda’)する

x = x.to(device)

# ネットワークへの正解ラベルyをy.to(‘cuda’)する

y = y.to(device)

これでGPU上のnetへxとyを入れることができるので演算可能になります。

 

速度比較

 

class SimpleNet(nn.Module):クラスとしてネットワークを作成します。
n_inputs:2
バッチ数:4
n_output:1
n_hidden:1024
hidden layer 4層をもつ全5層のネットワークです。
ネットワークのインスタンスを2つ作って
net_cpu = SimpleNet()
net_gpu = SimpleNet()
エポック数:1000でそれぞれ回してみましょう。
結果は
CPU training time: 9.972002267837524 seconds
GPU training time: 2.5578291416168213 seconds
ということで。GPUのほうが高速です。
しかし、n_hiddenを64にすると、
CPU training time: 0.5419738292694092 seconds
GPU training time: 2.506857395172119 seconds
となり、node数が少ない場合はcpuのほうが高速になります。
GPUならいつでも高速というわけではないことに注意しましょう。

スクリプト

 

次回はDDPGコードをGPU対応していきます。