INNER CODE UNIT · Python

load_feature_spec

NVIDIA/DeepLearningExamples · PyTorch/Recommendation/DLRM/dlrm/scripts/main.py:184

def load_feature_spec(flags):
    if flags.dataset_type == 'synthetic_gpu' and not flags.synthetic_dataset_use_feature_spec:
        num_numerical = flags.synthetic_dataset_numerical_features
        categorical_sizes = [int(s) for s in FLAGS.synthetic_dataset_table_sizes]
        return FeatureSpec.get_default_feature_spec(number_of_numerical_features=num_numerical,
                                                    categorical_feature_cardinalities=categorical_sizes)
    fspec_path = os.path.join(flags.dataset, flags.feature_spec)
    return FeatureSpec.from_yaml(fspec_path)


class CudaGraphWrapper:
    def __init__(self, model, train_step, parallelize,
                 zero_grad, cuda_graphs=False, warmup_steps=20):

        self.cuda_graphs = cuda_graphs
        self.warmup_iters = warmup_steps
        self.graph = None
        self.stream = None

View source record →