Skip to content
Snippets Groups Projects
swag.c 2.56 KiB
Newer Older
  • Learn to ignore specific revisions
  • Joseph Redmon's avatar
    Joseph Redmon committed
    #include "network.h"
    #include "detection_layer.h"
    #include "cost_layer.h"
    #include "utils.h"
    #include "parser.h"
    #include "box.h"
    
    #ifdef OPENCV
    #include "opencv2/highgui/highgui_c.h"
    #endif
    
    void train_swag(char *cfgfile, char *weightfile)
    {
        char *train_images = "data/voc.0712.trainval";
        char *backup_directory = "/home/pjreddie/backup/";
        srand(time(0));
        data_seed = time(0);
        char *base = basecfg(cfgfile);
        printf("%s\n", base);
        float avg_loss = -1;
        network net = parse_network_cfg(cfgfile);
        if(weightfile){
            load_weights(&net, weightfile);
        }
        printf("Learning Rate: %g, Momentum: %g, Decay: %g\n", net.learning_rate, net.momentum, net.decay);
        int imgs = net.batch*net.subdivisions;
        int i = *net.seen/imgs;
        data train, buffer;
    
        layer l = net.layers[net.n - 1];
    
        int side = l.side;
        int classes = l.classes;
        float jitter = l.jitter;
    
        list *plist = get_paths(train_images);
        //int N = plist->size;
        char **paths = (char **)list_to_array(plist);
    
        load_args args = {0};
        args.w = net.w;
        args.h = net.h;
        args.paths = paths;
        args.n = imgs;
        args.m = plist->size;
        args.classes = classes;
        args.jitter = jitter;
        args.num_boxes = side;
        args.d = &buffer;
        args.type = REGION_DATA;
    
        pthread_t load_thread = load_data_in_thread(args);
        clock_t time;
        //while(i*imgs < N*120){
        while(get_current_batch(net) < net.max_batches){
            i += 1;
            time=clock();
            pthread_join(load_thread, 0);
            train = buffer;
            load_thread = load_data_in_thread(args);
    
            printf("Loaded: %lf seconds\n", sec(clock()-time));
    
            time=clock();
            float loss = train_network(net, train);
            if (avg_loss < 0) avg_loss = loss;
            avg_loss = avg_loss*.9 + loss*.1;
    
            printf("%d: %f, %f avg, %f rate, %lf seconds, %d images\n", i, loss, avg_loss, get_current_rate(net), sec(clock()-time), i*imgs);
            if(i%1000==0 || i == 600){
                char buff[256];
                sprintf(buff, "%s/%s_%d.weights", backup_directory, base, i);
                save_weights(net, buff);
            }
            free_data(train);
        }
        char buff[256];
        sprintf(buff, "%s/%s_final.weights", backup_directory, base);
        save_weights(net, buff);
    }
    
    void run_swag(int argc, char **argv)
    {
        if(argc < 4){
            fprintf(stderr, "usage: %s %s [train/test/valid] [cfg] [weights (optional)]\n", argv[0], argv[1]);
            return;
        }
    
        char *cfg = argv[3];
        char *weights = (argc > 4) ? argv[4] : 0;
        if(0==strcmp(argv[2], "train")) train_swag(cfg, weights);
    }