INNER CODE UNIT · Python

train

vietnh1009/QuickDraw · train.py:41

def train(opt):
    if torch.cuda.is_available():
        torch.cuda.manual_seed(123)
    else:
        torch.manual_seed(123)
    training_params = {"batch_size": opt.batch_size,
                       "shuffle": True}

    test_params = {"batch_size": opt.batch_size,
                   "shuffle": False}

    output_file = open(opt.saved_path + os.sep + "logs.txt", "w")
    output_file.write("Model's parameters: {}".format(vars(opt)))

    training_set = MyDataset(opt.data_path, opt.total_images_per_class, opt.ratio, "train")
    training_generator = DataLoader(training_set, **training_params)
    print ("there are {} images for training phase".format(training_set.__len__()))
    test_set = MyDataset(opt.data_path, opt.total_images_per_class, opt.ratio, "test")

View source record →