Performs backward propagation from the loss and parameter update using asynchronous training. THIS IS A EXPERIMENTAL FUNCTION FOR RESEARCH PURPOSE: From the loss, it performs backward propagation to get the gradients and do the parameter update. It fuses the tensors smaller
(self, loss, threshold=2097152)
| 920 | self.opt.step() |
| 921 | |
| 922 | def backward_and_partial_update(self, loss, threshold=2097152): |
| 923 | """Performs backward propagation from the loss and parameter update using asynchronous training. |
| 924 | |
| 925 | THIS IS A EXPERIMENTAL FUNCTION FOR RESEARCH PURPOSE: |
| 926 | From the loss, it performs backward propagation to get the gradients and do the parameter |
| 927 | update. It fuses the tensors smaller than the threshold value to reduce network latency, |
| 928 | as well as performing asynchronous training where one parameter partition is all-reduced |
| 929 | per iteration. The size of the parameter partition depends on the threshold value. |
| 930 | |
| 931 | Args: |
| 932 | loss(Tensor): loss is the objective function of the deep learning model |
| 933 | optimization, e.g. for classification problem it can be the output of the |
| 934 | softmax_cross_entropy function. |
| 935 | threshold(int): threshold is a parameter to control performance in fusing |
| 936 | the tensors. For the tensors of sizes smaller than threshold, they are to |
| 937 | be accumulated and fused before the all reduce operation. For the tensors |
| 938 | of its size larger than the threshold value, they are to be reduced directly |
| 939 | without fusion. |
| 940 | |
| 941 | Attributes: |
| 942 | self.partial(int): A counter to determine which partition to perform all-reduce. |
| 943 | This counter resets to zero automatlly after an update cycle of the full parameter |
| 944 | set. |
| 945 | """ |
| 946 | if not hasattr(self, "partial"): |
| 947 | self.partial = 0 |
| 948 | self.partial += 1 |
| 949 | k = 0 |
| 950 | plist = [] |
| 951 | acc = 0 |
| 952 | tenlist = [] |
| 953 | reduced = [] |
| 954 | for p, g in autograd.backward(loss): |
| 955 | # every parameters update locally |
| 956 | self.opt.update(p, g) |
| 957 | # then do the partial parameter sychronization |
| 958 | if p.size() > threshold: |
| 959 | # larger than threshold -> reduced directly |
| 960 | # k is the partition number of the full gradient set |
| 961 | k += 1 |
| 962 | if (k == self.partial): |
| 963 | self.all_reduce(p.data) |
| 964 | reduced.append(p) |
| 965 | else: |
| 966 | # smaller than threshold -> accumulate |
| 967 | plist.append(p.data) |
| 968 | tenlist.append(p) |
| 969 | acc += p.size() |
| 970 | if (acc > threshold): |
| 971 | k += 1 |
| 972 | if (k == self.partial): |
| 973 | self.fused_all_reduce(plist, send=False) |
| 974 | self.fused_all_reduce(plist) |
| 975 | reduced = tenlist |
| 976 | acc = 0 |
| 977 | plist = [] |
| 978 | tenlist = [] |
| 979 | if plist: |
no test coverage detected