Returns a set of blobs that are needed for a complete check point. They are blobs for the first gpu and iteration blobs.
(model)
| 904 | |
| 905 | |
| 906 | def GetCheckpointParams(model): |
| 907 | ''' |
| 908 | Returns a set of blobs that are needed for a complete check point. |
| 909 | They are blobs for the first gpu and iteration blobs. |
| 910 | ''' |
| 911 | (all_blobs, _) = _ComputeBlobsToSync(model) |
| 912 | first_gpu_blobs = { |
| 913 | b |
| 914 | for b in all_blobs |
| 915 | if str(b) |
| 916 | .startswith("{}_{}/".format(model._device_prefix, model._devices[0])) |
| 917 | } |
| 918 | |
| 919 | # Add iteration blobs that do not have namescope separately, since |
| 920 | # it is important to checkpoint iteration counter |
| 921 | iteration_blobs = set() |
| 922 | for op in model.net.Proto().op: |
| 923 | if op.type == 'Iter' or op.type == 'AtomicIter': |
| 924 | if not op.output[0].startswith("{}_".format(model._device_prefix)): |
| 925 | iteration_blobs.add(op.output[0]) |
| 926 | |
| 927 | return first_gpu_blobs.union(iteration_blobs) |
| 928 | |
| 929 | |
| 930 | def FinalizeAfterCheckpoint(model, blobs=None, cpu_mode=False): |
nothing calls this directly
no test coverage detected
searching dependent graphs…