Ich bin ein Jax -Anfänger und jemand mit Jax hat mir gesagt, dass wenn wir wiederholt Anrufe zu einem Scan /
Loop (z. B. wenn diese selbst für Loop selbst eingewickelt werden), könnte es besser sein, die Schleife als für zu verlassen, anstatt sie in einen Scan zu konvertieren, da die für Loop vollständig abgeschlossene. Die Kosten für wiederholtes Ausführen dieser Schleife werden den Scan teurer als die für . Dies traf mich nicht sofort, als ich anfing, meinen Code zu schreiben, sondern machte Sinn, darüber nachzudenken.
Code: Select all
for i in range(num_train_steps): # Currently fixed to be a for loop
for j in range(num_env_steps): # Currently fixed to be a for loop
act()
def act():
for k in range(num_algo_iters): # Currently playing around with making this one either a scan or a for loop
jax.lax.scan(rollout_func) # Currently fixed to be a scan
Die einzige Schleife im obigen Code, die ich zwischen Scan getestet habe, und für war die k -schleife, und dann variierte ich die variable num_env_steps zu 1, 100, 1000 und 10000, um zu sehen, ob die Anzahl der mit dem Acting () (). (Die Tests wurden mit 5 Iterationen für die K für Loop- und 2 Iterationen für den innersten Scan durchgeführt, obwohl diese im Allgemeinen variabel sind, wenn dies wichtig ist.) Die für die -Version wurde also nie schneller für die Anzahl der Wiederholungen, die ich ausprobiert habe. Erhöhen Sie die Wiederholungen noch mehr durch Einstellen von num_env_steps auf 100k oder 1 Million machen es schneller oder können wir immer eine für durch ein Scan ersetzen? Ich habe diese Frage, weil ich mich frage, ob ich versuche, meinen Code zu überoptimieren, indem ich jeden für in einen Scan . Daumen, der mir helfen kann, mich zu entscheiden, wann ich für verwendet werden soll und wann ich scan verwenden soll, wenn ich nur an solchen Beschleunigungen interessiert bin?
wurde übrigens jitten.
Mobile version