예제 #1
0
파일: main.py 프로젝트: zw0610/tf-operator
def main(args):

    model_storage_type = args.model_storage_type
    if (model_storage_type == "local" or model_storage_type == "oss"):
        print("The storage type is " + model_storage_type)
    else:
        raise Exception("Only supports storage types like local and OSS")

    if args.job_type == "Predict":
        logging.info("starting the predict job")
        predict(args)

    elif args.job_type == "Train":
        logging.info("starting the train job")
        model = train(args)

        if model is not None:
            logging.info("finish the model training, and start to dump model ")
            model_path = args.model_path
            dump_model(model, model_storage_type, model_path, args)

    elif args.job_type == "All":
        logging.info("starting the train and predict job")

    logging.info("Finish distributed XGBoost job")
예제 #2
0
    def train(self, model, feature_name, model_dir, name):
        model.train()
        model_directory = os.path.join(model_dir)
        if not os.path.exists(model_directory):
            os.makedirs(model_directory)

        filename = os.path.join(
            model_directory,
            name + '_' + 'train' + '_' + feature_name + '.pkl')
        dump_model(model, filename)

        return model
예제 #3
0
def dump_models(name, f1_score, time_mark,
                data_pipeline=DataPipeline,
                learn_pipeline=LearningPipeline):
    path = '../dumps/history/%s__%s__%s__' % (name, time_mark, f1_score)
    dump_model(path + 'data.bin', data_pipeline)
    dump_model(path + 'learn.bin', learn_pipeline)
    path = '../dumps/%s__' % (name)
    dump_model(path + 'data.bin', data_pipeline)
    dump_model(path + 'learn.bin', learn_pipeline)
예제 #4
0
def main(args):

    if args.job_type == "Predict":
        logging.info("starting the predict job")
        predict(args)

    elif args.job_type == "Train":
        logging.info("starting the train job")
        model = train(args)

        if model is not None:
            logging.info("finish the model training, and start to dump model ")
            model_storage_type = args.model_storage_type
            model_path = args.model_path
            dump_model(model, model_storage_type, model_path, args)

    elif args.job_type == "All":
        logging.info("starting the train and predict job")

    logging.info("Finish distributed XGBoost job")
예제 #5
0
def train(epoch):
    print("begin train")
    model.train()
    criterion_c = nn.CrossEntropyLoss()
    criterion_a = nn.MultiLabelSoftMarginLoss()
    # due to a problem I commented out these lines
    #if cfg.ENABLE_TRIPLET_WITH_COSINE:
        #criterion_t = cfg.TripletMarginLossCosine()
    #else:
    criterion_t = nn.TripletMarginLoss()
    print("build loss")
    triplet_loader_iter = iter(triplet_loader)
    print("triplet_loader")
    triplet_type = 0
    if cfg.ENABLE_INSHOP_DATASET:
        triplet_in_shop_loader_iter = iter(triplet_in_shop_loader)
        print("in shop")
    for batch_idx, (data, target) in enumerate(train_loader):
        print("my train")
        if batch_idx % cfg.TEST_INTERVAL == 0:
            print("test")
            test()
        category=target['category']
        attribute=target['attribute']
        data, category, attribute = data.cpu(), category.cpu(), attribute.cpu()
        data, category, attribute = Variable(data), Variable(category),Variable(attribute)

        print("get data")
        optimizer.zero_grad()
        output1 = model(data)[0]
        output2 = model(data)[1]
        attribute = attribute.type(torch.FloatTensor)

        classification_loss = criterion_c(output1, category)
        attribute_loss = criterion_a(output2, attribute)
        if cfg.TRIPLET_WEIGHT:
            print("use triplet")
            if cfg.ENABLE_INSHOP_DATASET and random.random() < cfg.INSHOP_DATASET_PRECENT:
                triplet_type = 1
                try:
                    data_tri_list = next(triplet_in_shop_loader_iter)
                except StopIteration:
                    triplet_in_shop_loader_iter = iter(triplet_in_shop_loader)
                    data_tri_list = next(triplet_in_shop_loader_iter)
            else:
                triplet_type = 0
                try:
                    data_tri_list = next(triplet_loader_iter)
                except StopIteration:
                    triplet_loader_iter = iter(triplet_loader)
                    data_tri_list = next(triplet_loader_iter)
            triplet_batch_size = data_tri_list[0].shape[0]
            data_tri = torch.cat(data_tri_list, 0)
            data_tri = data_tri.cpu()
            data_tri = Variable(data_tri, requires_grad=True)
            feats = model(data_tri)[1]
            triplet_loss = criterion_t(
                feats[:triplet_batch_size],
                feats[triplet_batch_size:2 * triplet_batch_size],
                feats[2 * triplet_batch_size:]
            )
            loss = classification_loss + triplet_loss * cfg.TRIPLET_WEIGHT+attribute_loss
        else:
            loss = classification_loss + attribute_loss
        loss.backward()
        optimizer.step()
        if batch_idx % cfg.LOG_INTERVAL == 0:
            if cfg.TRIPLET_WEIGHT:
                print('Train Epoch: {} [{}/{} ({:.0f}%)]\tAll Loss: {:.4f}\t'
                      'Triple Loss({}): {:.4f}\tClassification Loss: {:.4f}'.format(
                    epoch, batch_idx * len(data), len(train_loader.dataset),
                    100. * batch_idx / len(train_loader), loss.data[0], triplet_type,
                    triplet_loss.data[0], classification_loss.data[0]))
            else:
                print('Train Epoch: {} [{}/{} ({:.0f}%)]\tClassification Loss: {:.4f}'.format(
                    epoch, batch_idx * len(data), len(train_loader.dataset),
                           100. * batch_idx / len(train_loader), float(loss.data)))
        if batch_idx and batch_idx % cfg.DUMP_INTERVAL == 0:
            print('Model saved to {}'.format(dump_model(model, epoch, batch_idx)))

    print('Model saved to {}'.format(dump_model(model, epoch)))
예제 #6
0
def train(epoch):
    model.train()
    criterion_c = nn.CrossEntropyLoss()
    if cfg.ENABLE_TRIPLET_WITH_COSINE:
        criterion_t = cfg.TripletMarginLossCosine()
    else:
        criterion_t = nn.TripletMarginLoss()
    triplet_loader_iter = iter(triplet_loader)
    triplet_type = 0
    if cfg.ENABLE_INSHOP_DATASET:
        triplet_in_shop_loader_iter = iter(triplet_in_shop_loader)
    for batch_idx, (data, target) in enumerate(train_loader):
        if batch_idx % cfg.TEST_INTERVAL == 0:
            test()
        data, target = data.cuda(cfg.GPU_ID), target.cuda(cfg.GPU_ID)
        data, target = Variable(data), Variable(target)
        optimizer.zero_grad()
        outputs = model(data)[0]
        classification_loss = criterion_c(outputs, target)
        if cfg.TRIPLET_WEIGHT:
            if cfg.ENABLE_INSHOP_DATASET and random.random(
            ) < cfg.INSHOP_DATASET_PRECENT:
                triplet_type = 1
                try:
                    data_tri_list = next(triplet_in_shop_loader_iter)
                except StopIteration:
                    triplet_in_shop_loader_iter = iter(triplet_in_shop_loader)
                    data_tri_list = next(triplet_in_shop_loader_iter)
            else:
                triplet_type = 0
                try:
                    data_tri_list = next(triplet_loader_iter)
                except StopIteration:
                    triplet_loader_iter = iter(triplet_loader)
                    data_tri_list = next(triplet_loader_iter)
            triplet_batch_size = data_tri_list[0].shape[0]
            data_tri = torch.cat(data_tri_list, 0)
            data_tri = data_tri.cuda(cfg.GPU_ID)
            data_tri = Variable(data_tri, requires_grad=True)
            feats = model(data_tri)[1]
            triplet_loss = criterion_t(
                feats[:triplet_batch_size],
                feats[triplet_batch_size:2 * triplet_batch_size],
                feats[2 * triplet_batch_size:])
            loss = classification_loss + triplet_loss * cfg.TRIPLET_WEIGHT
        else:
            loss = classification_loss
        loss.backward()
        optimizer.step()
        if batch_idx % cfg.LOG_INTERVAL == 0:
            if cfg.TRIPLET_WEIGHT:
                print('Train Epoch: {} [{}/{} ({:.0f}%)]\tAll Loss: {:.4f}\t'
                      'Triple Loss({}): {:.4f}\tClassification Loss: {:.4f}'.
                      format(epoch, batch_idx * len(data),
                             len(train_loader.dataset),
                             100. * batch_idx / len(train_loader),
                             loss.data[0], triplet_type, triplet_loss.data[0],
                             classification_loss.data[0]))
            else:
                print(
                    'Train Epoch: {} [{}/{} ({:.0f}%)]\tClassification Loss: {:.4f}'
                    .format(epoch, batch_idx * len(data),
                            len(train_loader.dataset),
                            100. * batch_idx / len(train_loader),
                            loss.data[0]))
        if batch_idx and batch_idx % cfg.DUMP_INTERVAL == 0:
            print('Model saved to {}'.format(
                dump_model(model, epoch, batch_idx)))

    print('Model saved to {}'.format(dump_model(model, epoch)))
예제 #7
0
def test_upload_model(model, model_path, args):

    return dump_model(model, type="local", model_path=model_path, args=args)