Data Parallel Model creates a net with ops in one device grouped together. This will interleave the ops so that each op for each device is next to each other in the net. Kind of like combining decks of cards. This ensures that progress is made along the critical path roughly concurr
(model)
| 1932 | |
| 1933 | |
| 1934 | def _InterleaveOps(model): |
| 1935 | ''' |
| 1936 | Data Parallel Model creates a net with ops in one device grouped together. |
| 1937 | This will interleave the ops so that each op for each device is next |
| 1938 | to each other in the net. Kind of like combining decks of cards. This |
| 1939 | ensures that progress is made along the critical path roughly concurrently |
| 1940 | for each device, which is important due to the extra intra-node |
| 1941 | synchronization required for multi-device batch normalization. |
| 1942 | ''' |
| 1943 | orig_ops = list(model.net.Proto().op) |
| 1944 | num_devices = len(model._devices) |
| 1945 | num_ops_per_dev = len(orig_ops) // num_devices |
| 1946 | assert num_devices * num_ops_per_dev == len(orig_ops), \ |
| 1947 | 'Number of ops per device in original net is not uniform' |
| 1948 | new_ops = [] |
| 1949 | ops = {d: [] for d in range(num_devices)} |
| 1950 | for op in orig_ops: |
| 1951 | ops[op.device_option.device_id].append(op) |
| 1952 | |
| 1953 | for j in range(num_ops_per_dev): |
| 1954 | tp = None |
| 1955 | for d in model._devices: |
| 1956 | if tp is None: |
| 1957 | tp = ops[d][j].type |
| 1958 | new_ops.append(ops[d][j]) |
| 1959 | # Sanity |
| 1960 | assert ops[d][j].type == tp, \ |
| 1961 | "Type mismatch {} / {}".format(tp, ops[d][j].type) |
| 1962 | |
| 1963 | del model.net.Proto().op[:] |
| 1964 | model.net.Proto().op.extend(new_ops) |
| 1965 | |
| 1966 | |
| 1967 | def _CPUInterDeviceBatchNormalization(model): |