def launch(self, dataset, on_epoch_end):
self.stop_train = False
# Load dataset
data_loader, data_loader_test = self.load_data(dataset)
# Initialize the model for this run
class_count = self.parameters.cfg["classes"]
if not dataset.has_bckgnd_class:
class_count = class_count + 1
model = models.faster_rcnn(
train_mode=True,
classes=class_count,
input_size=self.parameters.cfg["input_size"])
model.to(self.device)
# Optimizer
optimizer = self.init_optimizer(model)
# Train and evaluate
trained_model = self.train_model(model, data_loader, data_loader_test,
optimizer, on_epoch_end)
# Save model
if not os.path.isdir(self.parameters.cfg["output_folder"]):
os.mkdir(self.parameters.cfg["output_folder"])
if not self.parameters.cfg["output_folder"].endswith('/'):
self.parameters.cfg["output_folder"] += '/'
str_datetime = datetime.now().strftime("%d-%m-%YT%Hh%Mm%Ss")
model_folder = self.parameters.cfg["output_folder"] + str_datetime + "/"
if not os.path.isdir(model_folder):
os.mkdir(model_folder)
# .pth
if self.parameters.cfg["export_pth"]:
model_path = model_folder + self.parameters.cfg[
"model_name"] + ".pth"
ikutils.save_pth(trained_model, model_path)
# .onnx
if self.parameters.cfg["export_onnx"]:
model_path = model_folder + self.parameters.cfg[
"model_name"] + ".onnx"
input_shape = [
1, 3, self.parameters.cfg["input_size"],
self.parameters.cfg["input_size"]
]
ikutils.save_onnx(trained_model, input_shape, self.device,
model_path)
def launch(self, dataset_dir, on_epoch_end):
self.stop_train = False
# Load dataset
data_loaders = self.load_data(dataset_dir)
# Setup the loss function
loss = torch.nn.CrossEntropyLoss()
# Initialize the model for this run
model = models.resnet(
model_name=self.parameters.cfg["model_name"],
train_mode=True,
use_pretrained=self.parameters.cfg["use_pretrained"],
feature_extract=self.parameters.cfg["feature_extract"],
classes=self.parameters.cfg["classes"])
optimizer = self.init_optimizer(model)
# Train and evaluate
model_ft, hist = self.train_model(model, data_loaders, loss, optimizer,
on_epoch_end)
# Save model
if not os.path.isdir(self.parameters.cfg["output_folder"]):
os.mkdir(self.parameters.cfg["output_folder"])
if not self.parameters.cfg["output_folder"].endswith('/'):
self.parameters.cfg["output_folder"] += '/'
str_datetime = datetime.now().strftime("%d-%m-%YT%Hh%Mm%Ss")
model_folder = self.parameters.cfg["output_folder"] + str_datetime + "/"
if not os.path.isdir(model_folder):
os.mkdir(model_folder)
# .pth
if self.parameters.cfg["export_pth"]:
model_path = model_folder + self.parameters.cfg[
"model_name"] + ".pth"
utils.save_pth(model_ft, model_path)
# .onnx
if self.parameters.cfg["export_onnx"]:
model_path = model_folder + self.parameters.cfg[
"model_name"] + ".onnx"
input_shape = [
1, 3, self.parameters.cfg["input_size"],
self.parameters.cfg["input_size"]
]
utils.save_onnx(model, input_shape, self.device, model_path)