oneflow.cuda¶

ONEFLOW.CUDA¶

Licensed under the Apache License, Version 2.0 (the “License”); you may not use this file except in compliance with the License. You may obtain a copy of the License at

http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an “AS IS” BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License.

class oneflow.cuda.BoolTensor¶: The tensortype oneflow.cuda.BoolTensor is not available.

class oneflow.cuda.ByteTensor¶: The tensortype oneflow.cuda.ByteTensor is not available.

class oneflow.cuda.CharTensor¶: The tensortype oneflow.cuda.CharTensor is not available.

class oneflow.cuda.DoubleTensor¶: The tensortype oneflow.cuda.DoubleTensor is not available.

class oneflow.cuda.FloatTensor¶: The tensortype oneflow.cuda.FloatTensor is not available.

class oneflow.cuda.HalfTensor¶: The tensortype oneflow.cuda.HalfTensor is not available.

class oneflow.cuda.IntTensor¶: The tensortype oneflow.cuda.IntTensor is not available.

class oneflow.cuda.LongTensor¶: The tensortype oneflow.cuda.LongTensor is not available.

oneflow.cuda.current_device() → int¶: Returns local rank as device index.

oneflow.cuda.device_count() → int¶: Returns the number of GPUs available.

oneflow.cuda.empty_cache() → None¶: Releases all unoccupied cached memory currently held by the caching allocators of all OneFlow streams so those can be re-allocated in OneFlow streams or other GPU application and visible in nvidia-smi.

Note

empty_cache() may enable one stream to release memory and then freed memory can be used by another stream. It may also help reduce fragmentation of GPU memory in certain cases.

oneflow.cuda.is_available() → bool¶: Returns a bool indicating if CUDA is currently available.

oneflow.cuda.manual_seed(seed: int) → None¶

The documentation is referenced from: https://pytorch.org/docs/1.10/generated/torch.cuda.manual_seed.html.

Sets the seed for generating random numbers for the current GPU. It’s safe to call this function if CUDA is not available; in that case, it is silently ignored.

Parameters: seed (int) – The desired seed.

Warning

If you are working with a multi-GPU model, this function is insufficient to get determinism. To seed all GPUs, use manual_seed_all().

oneflow.cuda.manual_seed_all(seed) → None¶

The documentation is referenced from: https://pytorch.org/docs/1.10/generated/torch.cuda.manual_seed_all.html.

Sets the seed for generating random numbers on all GPUs. It’s safe to call this function if CUDA is not available; in that case, it is silently ignored.

Parameters: seed (int) – The desired seed.

oneflow.cuda.set_device(device: Union[oneflow._oneflow_internal.device, str, int]) → None¶

The documentation is referenced from: https://pytorch.org/docs/stable/generated/torch.cuda.set_device.html.

Sets the current device. Usage of this function is discouraged in favor of device. In most cases it’s better to use CUDA_VISIBLE_DEVICES environmental variable.

Parameters: device (flow.device or int) – selected device. This function is a no-op if this argument is negative.

oneflow.cuda.synchronize(device: Optional[Union[oneflow._oneflow_internal.device, str, int]] = None) → None¶

Waits for all kernels in all streams on a CUDA device to complete.

Note

In the eager mode of oneflow, all operations will be converted into instructions executed in the virtual machine, so in order to comply with the semantics of synchronization, this function will call the eager.Sync() function before the device is synchronized, which may affect the operations executed in other devices.

Parameters: device (flow.device or int, optional) – device for which to synchronize. It uses the current device, given by current_device(), if device is None (default).