<commit_before>from common import modelzoo
import mxnet as mx
import os
import logging
import math
import argparse

def train_imagenet(args):

    # arguments to change
    lr = args.lr
    lr_steps = [int(i) for i in args.lr_steps.split(',')]
    lr_factor = args.lr_factor
    weight_sparsity = [float(i) for i in args.weight_sparsity.split(',')]
    bias_sparsity = [float(i) for i in args.bias_sparsity.split(',')]
    switch_epoch = [int(i) for i in args.switch_epoch.split(',')]
    do_pruning = args.do_pruning
    gpus = [int(i) for i in args.gpus.split(',')]

    # fixed arguments
    num_epoch = 20
    batch_size = 512
    label_name = 'softmax_label'

    # create data iterators
    train = mx.io.ImageRecordIter(
        path_imgrec        = '/home/ubuntu/data/train_480_q90.rec',
        label_width        = 1,
        preprocess_threads = 4,
        batch_size         = batch_size,
        data_shape         = (3,224,224),
        label_name         = label_name,
        rand_crop          = False,
        rand_mirror        = False,
        mean_r             = 0.0,
        mean_g             = 0.0,
        mean_b             = 0.0)
    val = mx.io.ImageRecordIter(
        path_imgrec        = '/home/ubuntu/data/val_480_q90.rec',
        label_width        = 1,
        preprocess_threads = 4,
        batch_size         = batch_size,
        data_shape         = (3,224,224),
        label_name         = label_name,
        rand_crop          = False,
        rand_mirror        = False,
        mean_r             = 0.0,
        mean_g             = 0.0,
        mean_b             = 0.0)

    # download model
    model = 'imagenet1k-resnet-18'
    dir_path = os.path.dirname(os.path.realpath(__file__))
    (prefix, epoch) = modelzoo.download_model(model, os.path.join(dir_path, 'model'))
    sym, arg_params, aux_params = mx.model.load_checkpoint(prefix, epoch)

    # training
    context = [mx.gpu(i) for i in gpus]
    mod = mx.mod.Module(symbol = sym, context = context, label_names = [label_name,])
    begin_epoch = 0
    eval_metric = ['accuracy']
    kv = 'local'
    optimizer = 'sgd'
    batches_per_epoch = math.ceil(1281167.0 / batch_size)
    lr_steps = [batches_per_epoch * i for i in lr_steps]
    lr_scheduler = mx.lr_scheduler.MultiFactorScheduler(step = lr_steps,
        factor = lr_factor)
    start_pruning = True
    optimizer_params = {
            'learning_rate'     : lr,
            'lr_scheduler'      : lr_scheduler,
            'weight_sparsity'   : weight_sparsity,
            'bias_sparsity'     : bias_sparsity,
            'switch_epoch'      : switch_epoch,
            'batches_per_epoch' : batches_per_epoch,
            'do_pruning'        : do_pruning,
            'start_prune'       : start_pruning,
    }
    disp_batches = 50
    batch_end_callbacks = [mx.callback.Speedometer(batch_size, disp_batches)]
    mod.fit(train,
        begin_epoch                 = begin_epoch,
        num_epoch                   = num_epoch,
        eval_data                   = val,<commit_msg>changes to imagenet pruning<commit_after>from common import modelzoo
import mxnet as mx
import os
import logging
import math
import argparse

def train_imagenet(args):

    # arguments to change
    lr = args.lr
    lr_steps = [int(i) for i in args.lr_steps.split(',')]
    lr_factor = args.lr_factor
    weight_sparsity = [float(i) for i in args.weight_sparsity.split(',')]
    bias_sparsity = [float(i) for i in args.bias_sparsity.split(',')]
    switch_epoch = [int(i) for i in args.switch_epoch.split(',')]
    do_pruning = args.do_pruning
    gpus = [int(i) for i in args.gpus.split(',')]

    # fixed arguments
    num_epoch = 40
    batch_size = 128
    label_name = 'softmax_label'

    # create data iterators
    """
    train = mx.io.ImageRecordIter(
        path_imgrec        = '/home/ubuntu/data/train_480_q90.rec',
        label_width        = 1,
        preprocess_threads = 4,
        batch_size         = batch_size,
        data_shape         = (3,224,224),
        label_name         = label_name,
        rand_crop          = False,
        rand_mirror        = False,
        mean_r             = 0.0,
        mean_g             = 0.0,
        mean_b             = 0.0)
    val = mx.io.ImageRecordIter(
        path_imgrec        = '/home/ubuntu/data/val_480_q90.rec',
        label_width        = 1,
        preprocess_threads = 4,
        batch_size         = batch_size,
        data_shape         = (3,224,224),
        label_name         = label_name,
        rand_crop          = False,
        rand_mirror        = False,
        mean_r             = 0.0,
        mean_g             = 0.0,
        mean_b             = 0.0)
    """
    train = mx.io.ImageRecordIter(
        path_imgrec         = '/home/ubuntu/data/train_480_q90.rec',
        label_width         = 1,
        data_name           = 'data',
        label_name          = label_name,
        data_shape          = (3, 224, 224),
        batch_size          = batch_size,
        pad                 = 0,
        fill_value          = 127, # only used when pad is valid
        rand_crop           = True,
        max_random_scale    = 1.0, # 480 with imagnet, 32 with cifar10
        min_random_scale    = 0.533, # 256.0/480.0
        max_aspect_ratio    = 0.25,
        random_h            = 36, # 0.4*90
        random_s            = 50, # 0.4*127
        random_l            = 50, # 0.4*127
        rand_mirror         = True,
        shuffle             = True)
    val = mx.io.ImageRecordIter(
        path_imgrec         = '/home/ubuntu/data/val_256_q90.rec',
        max_random_scale    = 1, #0.533, # 480 with imagnet, 32 with cifar10
        min_random_scale    = 1, #0.533, # 256.0/480.0
        label_width         = 1,
        data_name           = 'data',
        label_name          = 'softmax_label',
        batch_size          = batch_size,
        data_shape          = (3, 224, 224),
        rand_crop           = False,
        rand_mirror         = False)

    # download model
    model = 'imagenet1k-resnet-101'
    dir_path = os.path.dirname(os.path.realpath(__file__))
    (prefix, epoch) = modelzoo.download_model(model, os.path.join(dir_path, 'model'))
    sym, arg_params, aux_params = mx.model.load_checkpoint(prefix, epoch)

    # training
    context = [mx.gpu(i) for i in gpus]
    mod = mx.mod.Module(symbol = sym, context = context, label_names = [label_name,])
    begin_epoch = 0
    eval_metric = ['accuracy']
    kv = 'local'
    optimizer = 'sgd'
    batches_per_epoch = math.ceil(1281167.0 / batch_size)
    lr_steps = [batches_per_epoch * i for i in lr_steps]
    lr_scheduler = mx.lr_scheduler.MultiFactorScheduler(step = lr_steps,
        factor = lr_factor)
    start_pruning = True
    optimizer_params = {
            'learning_rate'     : lr,
            'lr_scheduler'      : lr_scheduler,
            'weight_sparsity'   : weight_sparsity,
            'bias_sparsity'     : bias_sparsity,
            'switch_epoch'      : switch_epoch,
            'batches_per_epoch' : batches_per_epoch,
            'do_pruning'        : do_pruning,
            'start_prune'       : start_pruning,
    }
    disp_batches = 5
    batch_end_callbacks = [mx.callback.Speedometer(batch_size, disp_batches)]
    mod.fit(train,
        begin_epoch                 = begin_epoch,
        num_epoch                   = num_epoch,
        eval_data                   = val,
{'fraction_non_alphanumeric': 0.09814715359828141, 'fraction_numerical': 0.02913533834586466, 'mean_word_length': 1.733577981651376, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5891919', 'n_tokens_mistral': 2633, 'n_tokens_neox': 2498, 'n_words': 489}