Automatic GPU selection

[jprokos26]

Currently when gpu is not specified, armory defaults to running the attack on GPU0 and loading some other data onto all other GPUs. It would be useful to be able to launch armory without having to check gpustat.

When gpu is not set to a specific integer value and use_gpu is set to True, utilize the following function to automatically detect the best gpu to use as done elsewhere:

def get_device(device=None) -> torch.device:
    """
    If None is passed, the GPU with the most free memory will be used.
    If False is passed, the CPU will be used.
    Otherwise, the device will be set to the specified GPU.
    """
    if device is False:
        log.info("Using CPU")
        return torch.device("cpu")
    elif device is None:
        # Get the current memory usage of all GPUs using nvidia-smi
        memory = {
            i: int(
                os.popen(
                    f"nvidia-smi --query-gpu=memory.free --format=csv,noheader,nounits --id={i}"
                ).read()
            )
            for i in range(torch.cuda.device_count())
        }

        # Get the GPU with the most free memory
        device = torch.device(f"cuda:{max(memory, key=memory.get)}")
        log.info(
            f"Autoselected {device} with {memory[device.index]} MB bytes free"
        )
        return device
    else:
        try:
            return torch.device(f"cuda:{device}")
        except Exception as e:
            log.error(f"Error setting device to {device}: {e}")
            return torch.device("cpu")

[swsuggs]

@christopherwoodall this was not completed by 1925, is this no longer desired?