use_cpu=args.use_cpu) # for cifar10_quick_train_test.prototxt # train((train_x, train_y, test_x, test_y), net, 18, caffe_lr, 0.004, # use_cpu=args.use_cpu) elif args.model == 'alexnet': train_x, test_x = normalize_for_alexnet(train_x, test_x) net = alexnet.create_net(args.use_cpu) train((train_x, train_y, test_x, test_y), net, 2, alexnet_lr, 0.004, use_cpu=args.use_cpu) elif args.model == 'vgg': train_x, test_x = normalize_for_vgg(train_x, test_x) net = vgg.create_net(args.use_cpu) train((train_x, train_y, test_x, test_y), net, 250, vgg_lr, 0.0005, use_cpu=args.use_cpu) else: train_x, test_x = normalize_for_alexnet(train_x, test_x) net = resnet.create_net(args.use_cpu) train((train_x, train_y, test_x, test_y), net, 200, resnet_lr, 1e-4, use_cpu=args.use_cpu)
args = parser.parse_args() assert os.path.exists(args.data), \ 'Pls download the cifar10 dataset via "download_data.py py"' print('Loading data ..................') train_x, train_y = load_train_data(args.data) test_x, test_y = load_test_data(args.data) if args.model == 'caffe': train_x, test_x = normalize_for_alexnet(train_x, test_x) net = caffe_net.create_net(args.use_cpu) # for cifar10_full_train_test.prototxt train((train_x, train_y, test_x, test_y), net, 160, alexnet_lr, 0.004, use_cpu=args.use_cpu) # for cifar10_quick_train_test.prototxt # train((train_x, train_y, test_x, test_y), net, 18, caffe_lr, 0.004, # use_cpu=args.use_cpu) elif args.model == 'cnn': train_x, test_x = normalize_for_alexnet(train_x, test_x) net = cnn.create_net(args.use_cpu) train((train_x, train_y, test_x, test_y), net, 2, alexnet_lr, 0.004, use_cpu=args.use_cpu) elif args.model == 'vgg': train_x, test_x = normalize_for_vgg(train_x, test_x) net = vgg.create_net(args.use_cpu) train((train_x, train_y, test_x, test_y), net, 250, vgg_lr, 0.0005, use_cpu=args.use_cpu) else: train_x, test_x = normalize_for_alexnet(train_x, test_x) net = resnet.create_net(args.use_cpu) train((train_x, train_y, test_x, test_y), net, 200, resnet_lr, 1e-4, use_cpu=args.use_cpu)