run.py

import argparse
import logging
import os
import subprocess
from io import StringIO

import pandas as pd

import configs
import utils
from datahandler import DataHandler
from dgc import DGC


def main(tag, seed, dataset):
    opts = getattr(configs, 'config_%s' % dataset)
    opts['work_dir'] = './results/%s/' % tag

    if opts['verbose']:
        logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(message)s')
    logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(message)s')
    utils.create_dir(opts['work_dir'])
    utils.create_dir(os.path.join(opts['work_dir'],
                                  'checkpoints'))

    with utils.o_gfile((opts['work_dir'], 'params.txt'), 'w') as text:
        text.write('Parameters:\n')
        for key in opts:
            text.write('%s : %s\n' % (key, opts[key]))

    data = DataHandler(opts, seed)
    model = DGC(opts, tag)
    model.train(data)


def get_free_gpu(num=1):
    gpu_stats = subprocess.check_output(["nvidia-smi", "--format=csv", "--query-gpu=memory.used,memory.free"])
    gpu_df = pd.read_csv(StringIO(gpu_stats.decode('utf8')),
                         names=['memory.used', 'memory.free'],
                         skiprows=1)
    gpu_df['memory.free'] = gpu_df['memory.free'].map(lambda x: int(x.rstrip(' [MiB]')))
    gpu_df = gpu_df.sort_values(by='memory.free', ascending=False)
    print('GPU usage:\n{}'.format(gpu_df))
    free_gpus = []
    for i in range(num):
        print('Returning GPU{} with {} free MiB'.format(gpu_df.index[i], gpu_df.iloc[i]['memory.free']))
        free_gpus.append(str(gpu_df.index[i]))
    return ','.join(free_gpus)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("--exp", default='mnist',
                        help='dataset [mnist/celeba]')
    parser.add_argument("--seed", type=int, default=1,
                        help='random seed for imbalance data generation')
    FLAGS = parser.parse_args()

    os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
    free_gpu_id = get_free_gpu(num=1)
    os.environ["CUDA_VISIBLE_DEVICES"] = free_gpu_id
    os.environ["OMP_NUM_THREADS"] = "8"

    dataset_name = FLAGS.exp
    seed = FLAGS.seed
    tag = '%s_seed%02d' % (dataset_name, seed)
    main(tag, seed, dataset_name)