Spaces:
Sleeping
Sleeping
| import time | |
| import torch | |
| import torch.nn.functional as F | |
| from hpc_rll.origin.vtrace import vtrace_error_discrete_action, vtrace_data | |
| from hpc_rll.rl_utils.vtrace import VTrace | |
| from testbase import mean_relative_error, times | |
| assert torch.cuda.is_available() | |
| use_cuda = True | |
| T = 128 | |
| B = 128 | |
| N = 128 | |
| def vtrace_val(): | |
| ori_target_output = torch.randn(T, B, N) | |
| ori_behaviour_output = torch.randn(T, B, N) | |
| ori_action = torch.randint( | |
| 0, N, size=( | |
| T, | |
| B, | |
| ) | |
| ) | |
| ori_value = torch.randn(T + 1, B) | |
| ori_reward = torch.randn(T, B) | |
| hpc_target_output = ori_target_output.clone().detach() | |
| hpc_behaviour_output = ori_behaviour_output.clone().detach() | |
| hpc_action = ori_action.clone().detach() | |
| hpc_value = ori_value.clone().detach() | |
| hpc_reward = ori_reward.clone().detach() | |
| hpc_vtrace = VTrace(T, B, N) | |
| if use_cuda: | |
| ori_target_output = ori_target_output.cuda() | |
| ori_behaviour_output = ori_behaviour_output.cuda() | |
| ori_action = ori_action.cuda() | |
| ori_value = ori_value.cuda() | |
| ori_reward = ori_reward.cuda() | |
| hpc_target_output = hpc_target_output.cuda() | |
| hpc_behaviour_output = hpc_behaviour_output.cuda() | |
| hpc_action = hpc_action.cuda() | |
| hpc_value = hpc_value.cuda() | |
| hpc_reward = hpc_reward.cuda() | |
| hpc_vtrace = hpc_vtrace.cuda() | |
| ori_target_output.requires_grad_(True) | |
| ori_value.requires_grad_(True) | |
| ori_loss = vtrace_error_discrete_action( | |
| vtrace_data(ori_target_output, ori_behaviour_output, ori_action, ori_value, ori_reward, None) | |
| ) | |
| ori_loss = sum(ori_loss) | |
| ori_loss.backward() | |
| hpc_target_output.requires_grad_(True) | |
| hpc_value.requires_grad_(True) | |
| hpc_loss = hpc_vtrace(hpc_target_output, hpc_behaviour_output, hpc_action, hpc_value, hpc_reward) | |
| hpc_loss = sum(hpc_loss) | |
| hpc_loss.backward() | |
| mre = mean_relative_error( | |
| torch.flatten(ori_loss).cpu().detach().numpy(), | |
| torch.flatten(hpc_loss).cpu().detach().numpy() | |
| ) | |
| print("vtrace fp mean_relative_error: " + str(mre)) | |
| mre = mean_relative_error( | |
| torch.flatten(ori_target_output.grad).cpu().detach().numpy(), | |
| torch.flatten(hpc_target_output.grad).cpu().detach().numpy() | |
| ) | |
| print("vtrace bp target_output mean_relative_error: " + str(mre)) | |
| mre = mean_relative_error( | |
| torch.flatten(ori_value.grad).cpu().detach().numpy(), | |
| torch.flatten(hpc_value.grad).cpu().detach().numpy() | |
| ) | |
| print("vtrace bp value mean_relative_error: " + str(mre)) | |
| def vtrace_perf(): | |
| ori_target_output = torch.randn(T, B, N) | |
| ori_behaviour_output = torch.randn(T, B, N) | |
| ori_action = torch.randint( | |
| 0, N, size=( | |
| T, | |
| B, | |
| ) | |
| ) | |
| ori_value = torch.randn(T + 1, B) | |
| ori_reward = torch.randn(T, B) | |
| hpc_target_output = ori_target_output.clone().detach() | |
| hpc_behaviour_output = ori_behaviour_output.clone().detach() | |
| hpc_action = ori_action.clone().detach() | |
| hpc_value = ori_value.clone().detach() | |
| hpc_reward = ori_reward.clone().detach() | |
| hpc_vtrace = VTrace(T, B, N) | |
| if use_cuda: | |
| ori_target_output = ori_target_output.cuda() | |
| ori_behaviour_output = ori_behaviour_output.cuda() | |
| ori_action = ori_action.cuda() | |
| ori_value = ori_value.cuda() | |
| ori_reward = ori_reward.cuda() | |
| hpc_target_output = hpc_target_output.cuda() | |
| hpc_behaviour_output = hpc_behaviour_output.cuda() | |
| hpc_action = hpc_action.cuda() | |
| hpc_value = hpc_value.cuda() | |
| hpc_reward = hpc_reward.cuda() | |
| hpc_vtrace = hpc_vtrace.cuda() | |
| ori_target_output.requires_grad_(True) | |
| ori_value.requires_grad_(True) | |
| for i in range(times): | |
| t = time.time() | |
| ori_loss = vtrace_error_discrete_action( | |
| vtrace_data(ori_target_output, ori_behaviour_output, ori_action, ori_value, ori_reward, None) | |
| ) | |
| ori_loss = sum(ori_loss) | |
| ori_loss.backward() | |
| if use_cuda: | |
| torch.cuda.synchronize() | |
| print('epoch: {}, original vtrace cost time: {}'.format(i, time.time() - t)) | |
| hpc_target_output.requires_grad_(True) | |
| hpc_value.requires_grad_(True) | |
| for i in range(times): | |
| t = time.time() | |
| hpc_loss = hpc_vtrace(hpc_target_output, hpc_behaviour_output, hpc_action, hpc_value, hpc_reward) | |
| hpc_loss = sum(hpc_loss) | |
| hpc_loss.backward() | |
| if use_cuda: | |
| torch.cuda.synchronize() | |
| print('epoch: {}, hpc vtrace cost time: {}'.format(i, time.time() - t)) | |
| if __name__ == '__main__': | |
| print("target problem: T = {}, B = {}, N = {}".format(T, B, N)) | |
| print("================run vtrace validation test================") | |
| vtrace_val() | |
| print("================run vtrace performance test================") | |
| vtrace_perf() | |