diff --git a/tests/unit_tests/test_optimizer.py b/tests/unit_tests/test_optimizer.py index 776e3c2..291e30a 100644 --- a/tests/unit_tests/test_optimizer.py +++ b/tests/unit_tests/test_optimizer.py @@ -1,3 +1,4 @@ +import os import autogen from opto.trace import bundle, node, GRAPH from opto.optimizers import OptoPrime @@ -123,60 +124,62 @@ def foobar_text(x): GRAPH.clear() x = node("negative point one", trainable=True) -optimizer = OptoPrime([x], config_list=autogen.config_list_from_json("OAI_CONFIG_LIST")) -output = foobar_text(x) -feedback = user(output.data) -optimizer.zero_feedback() -optimizer.backward(output, feedback) -print(f"variable={x.data}, output={output.data}, feedback={feedback}") # logging -optimizer.step(verbose=True) - -## Test the optimizer with an example of code -GRAPH.clear() - - -def user(output): - if output < 0: - return "Success." - else: - return "Try again. The output should be negative" - - -# We make this function as a parameter that can be optimized. -@bundle(trainable=True) -def my_fun(x): - """Test function""" - return x**2 + 1 - -x = node(-1, trainable=False) -optimizer = OptoPrime([my_fun.parameter], config_list=autogen.config_list_from_json("OAI_CONFIG_LIST")) -output = my_fun(x) -feedback = user(output.data) -optimizer.zero_feedback() -optimizer.backward(output, feedback) - -print(f"output={output.data}, feedback={feedback}, variables=\n") # logging -for p in optimizer.parameters: - print(p.name, p.data) -optimizer.step(verbose=True) - - -# Test directly providing feedback to parameters -GRAPH.clear() -x = node(-1, trainable=True) -optimizer = OptoPrime([x]) -feedback = "test" -optimizer.zero_feedback() -optimizer.backward(x, feedback) -optimizer.step(verbose=True) - - -# Test if we can save log in both pickle and json -import json, pickle -json.dump(optimizer.log, open("log.json", "w")) -pickle.dump(optimizer.log, open("log.pik", "wb")) -# remove these files -import os -os.remove("log.json") -os.remove("log.pik") \ No newline at end of file +if os.path.exists("OAI_CONFIG_LIST"): + optimizer = OptoPrime([x], config_list=autogen.config_list_from_json("OAI_CONFIG_LIST")) + output = foobar_text(x) + feedback = user(output.data) + optimizer.zero_feedback() + optimizer.backward(output, feedback) + print(f"variable={x.data}, output={output.data}, feedback={feedback}") # logging + optimizer.step(verbose=True) + + ## Test the optimizer with an example of code + GRAPH.clear() + + + def user(output): + if output < 0: + return "Success." + else: + return "Try again. The output should be negative" + + + # We make this function as a parameter that can be optimized. + @bundle(trainable=True) + def my_fun(x): + """Test function""" + return x**2 + 1 + + + x = node(-1, trainable=False) + optimizer = OptoPrime([my_fun.parameter], config_list=autogen.config_list_from_json("OAI_CONFIG_LIST")) + output = my_fun(x) + feedback = user(output.data) + optimizer.zero_feedback() + optimizer.backward(output, feedback) + + print(f"output={output.data}, feedback={feedback}, variables=\n") # logging + for p in optimizer.parameters: + print(p.name, p.data) + optimizer.step(verbose=True) + + + # Test directly providing feedback to parameters + GRAPH.clear() + x = node(-1, trainable=True) + optimizer = OptoPrime([x]) + feedback = "test" + optimizer.zero_feedback() + optimizer.backward(x, feedback) + optimizer.step(verbose=True) + + + # Test if we can save log in both pickle and json + import json, pickle + json.dump(optimizer.log, open("log.json", "w")) + pickle.dump(optimizer.log, open("log.pik", "wb")) + # remove these files + import os + os.remove("log.json") + os.remove("log.pik") \ No newline at end of file