Method Train

lesson7-Detection/src/Detector.cpp:98–233

}

void Detector::Train(std::string train_val_path, std::string image_type, int num_epochs, int batch_size,
	float learning_rate, std::string save_path, std::string pretrained_path) {
	if (!does_exist(pretrained_path))
	{
		std::cout << "Pretrained path is invalid: " << pretrained_path <<"\t random initialzed the model"<< std::endl;
	}
	else {
		loadPretrained(pretrained_path);
	}
	std::string train_label_path = train_val_path + "\\train\\images";
	std::string train_image_path = train_val_path + "\\train\\labels";
	std::string val_label_path = train_val_path + "\\val\\images";
	std::string val_image_path = train_val_path + "\\val\\labels";

	std::vector<std::string> list_images_train = {};
	std::vector<std::string> list_labels_train = {};
	std::vector<std::string> list_images_val = {};
	std::vector<std::string> list_labels_val = {};

	load_det_data_from_folder(train_image_path, image_type, list_images_train, list_labels_train);
	load_det_data_from_folder(val_image_path, image_type, list_images_val, list_labels_val);

	if (list_images_train.size() < batch_size || list_images_val.size() < batch_size) {
		std::cout << "Image numbers less than batch size or empty image folder" << std::endl;
		return;
	}
	if (!does_exist(list_images_train[0]))
	{
		std::cout << "Image path is invalid get first train image " << list_images_train[0] << std::endl;
		return;
	}
	auto custom_dataset_train = DetDataset(list_images_train, list_labels_train, name_list, true,
		width, height);
	auto custom_dataset_val = DetDataset(list_images_val, list_labels_val, name_list, false, width, height);
	auto data_loader_train = torch::data::make_data_loader<torch::data::samplers::RandomSampler>(std::move(custom_dataset_train), batch_size);
	auto data_loader_val = torch::data::make_data_loader<torch::data::samplers::RandomSampler>(std::move(custom_dataset_val), batch_size);

	float anchor[12] = { 10,14,  23,27,  37,58,  81,82,  135,169,  344,319 };
	auto anchors_ = torch::from_blob(anchor, { 6,2 }, torch::TensorOptions(torch::kFloat32)).to(device);
	int image_size[2] = { width, height };

	bool normalize = false;
	auto critia1 = YOLOLossImpl(anchors_, name_list.size(), image_size, 0.01, device, normalize);
	auto critia2 = YOLOLossImpl(anchors_, name_list.size(), image_size, 0.01, device, normalize);
	
	auto pretrained_dict = detector->named_parameters();
	auto FloatType = torch::ones(1).to(torch::kFloat).to(device).options();
	for (int epoc_count = 0; epoc_count < num_epochs; epoc_count++) {
		float loss_sum = 0;
		int batch_count = 0;
		float loss_train = 0;
		float loss_val = 0;
		float best_loss = 1e10;

		if (epoc_count == int(num_epochs / 2)) { learning_rate /= 10; }
		torch::optim::Adam optimizer(detector->parameters(), learning_rate); // Learning Rate
		if (epoc_count < int(num_epochs / 10)) {
			for (auto mm : pretrained_dict)