initialize with the necessary elements
(self, use_gpu=False, enable_mkldnn=False)
| 47 | ) |
| 48 | class TableSystem(hub.Module): |
| 49 | def _initialize(self, use_gpu=False, enable_mkldnn=False): |
| 50 | """ |
| 51 | initialize with the necessary elements |
| 52 | """ |
| 53 | cfg = self.merge_configs() |
| 54 | cfg.use_gpu = use_gpu |
| 55 | if use_gpu: |
| 56 | try: |
| 57 | _places = os.environ["CUDA_VISIBLE_DEVICES"] |
| 58 | int(_places[0]) |
| 59 | print("use gpu: ", use_gpu) |
| 60 | print("CUDA_VISIBLE_DEVICES: ", _places) |
| 61 | cfg.gpu_mem = 8000 |
| 62 | except: |
| 63 | raise RuntimeError( |
| 64 | "Environment Variable CUDA_VISIBLE_DEVICES is not set correctly. If you wanna use gpu, please set CUDA_VISIBLE_DEVICES via export CUDA_VISIBLE_DEVICES=cuda_device_id." |
| 65 | ) |
| 66 | cfg.ir_optim = True |
| 67 | cfg.enable_mkldnn = enable_mkldnn |
| 68 | |
| 69 | self.table_sys = _TableSystem(cfg) |
| 70 | |
| 71 | def merge_configs(self): |
| 72 | # default cfg |