def on_init(self, **kwargs):
mode = kwargs.get('mode', None)
input_transform_dir = kwargs.get('input_image_transformation', None)
self.transform = self.get_transforms(input_transform_dir, mode)
logger.info(f'{mode}, transforms {self.transform}.')
logger.info("Transformation init finished.")
def run(self, **kwargs):
input_image_dir = kwargs.get('input_image_directory', None)
logger.info(f'Applying transform:')
transformed_dir = input_image_dir.apply_to_images(
transform=lambda image: self.unloader(
self.transform(image).squeeze(0)).convert("RGB"))
return (transformed_dir, )
def entrance(input_path='/mnt/chjinche/data/small/',
output_path='/mnt/chjinche/data/image_dir/'):
logger.info('Start!')
# Case 1: input path is torchvision ImageFolder
loader_dir = FolderBasedImageDirectory.load_organized(input_path)
loader_dir.dump(output_path)
logger.info('Finished')
def apply(self, ports={}, params={}):
input_image_dir = ports.get('Input image directory', None)
logger.info(f'Applying transform:')
transformed_dir = input_image_dir.apply_to_images(
transform=lambda image: self.unloader(
self.transform(image).squeeze(0)))
return transformed_dir
def onstart(self, ports={}, params={}):
mode = params.get('Mode', None)
input_transform_dir = ports.get('Input image transformation', None)
self.transform = self.get_transforms(input_transform_dir, mode)
logger.info(f'{mode}, transforms {self.transform}.')
logger.info("Transformation init finished.")
def entrance(input_path='../image_dataset/', output_path='../image_dir/'):
logger.info('Start!')
# Case 1: input path is torchvision ImageFolder
# TODO: Case 2: input path is custom image format
loader_dir = FolderBasedImageDirectory.load_organized(input_path)
loader_dir.dump(output_path)
logger.info('Finished')
def entrance(src_path='../image_dir/',
fraction=0.9,
tgt_train_path='../image_dir_train/',
tgt_test_path='../image_dir_test/'):
logger.info('Start!')
split_images(src_path, tgt_train_path, tgt_test_path, fraction)
logger.info('Finished')
def entrance(
src_path='/mnt/chjinche/test_data/detection/image_dir/',
fraction=0.9,
tgt_train_path='/mnt/chjinche/test_data/detection/image_dir_train/',
tgt_test_path='/mnt/chjinche/test_data/detection/image_dir_test/'):
logger.info('Start!')
split_images(src_path, tgt_train_path, tgt_test_path, fraction)
logger.info('Finished')
def __init__(self, input_transform_path, transform_type):
self.transform = self.get_transforms(input_transform_path,
transform_type)
logger.info(
f'Set transform_type {transform_type}, transforms {self.transform}.'
)
self.unloader = transforms.ToPILImage()
logger.info("Transformation init finished.")
def entrance(resize_size=256,
center_crop_size=224,
five_crop=False,
ten_crop=False,
pad=False,
color_jitter=False,
grayscale=False,
random_crop=False,
random_horizontal_flip=True,
random_vertical_flip=False,
random_resized_crop=False,
random_rotation=False,
random_affine=False,
random_grayscale=False,
random_perspective=False,
random_erasing=False,
normalize=True,
output_path='../init_transform/'):
# Construct image transform
# TODO: check transforms ordering
img_trans_dir = ImageTransformationDirectory.create(
transforms=[('Resize', resize_size), ('CenterCrop', center_crop_size)])
if five_crop:
img_trans_dir.append('FiveCrop')
if ten_crop:
img_trans_dir.append('TenCrop')
if pad:
img_trans_dir.append('Pad')
if color_jitter:
img_trans_dir.append('ColorJitter')
if grayscale:
img_trans_dir.append('Grayscale')
if random_crop:
img_trans_dir.append('RandomCrop')
if random_horizontal_flip:
img_trans_dir.append('RandomHorizontalFlip')
if random_vertical_flip:
img_trans_dir.append('RandomVerticalFlip')
if random_resized_crop:
img_trans_dir.append('RandomResizedCrop')
if random_rotation:
img_trans_dir.append('RandomRotation')
if random_affine:
img_trans_dir.append('RandomAffine')
if random_grayscale:
img_trans_dir.append('RandomGrayscale')
if random_perspective:
img_trans_dir.append('RandomPerspective')
if random_erasing:
img_trans_dir.append('RandomErasing')
# Need to do 'ToTensor' op ahead of normalzation.
img_trans_dir.append('ToTensor')
if normalize:
img_trans_dir.append_normalize()
logger.info(f'Constructed image transforms: {img_trans_dir.transforms}')
# Dump
img_trans_dir.dump(output_path)
logger.info('Finished')
def entrance(input_path='../image_dataset/', output_path='../image_dir/'):
logger.info('Start!')
with TimeProfile(f"Mount/Download dataset to {input_path}"):
print_dir_hierarchy_to_log(input_path)
# Case 1: input path is torchvision ImageFolder
# TODO: Case 2: input path is custom image format
loader_dir = FolderBasedImageDirectory.load_organized(input_path)
loader_dir.dump(output_path)
logger.info('Finished.')
def get_split_list(ann_type, lst, fraction):
n = len(lst)
if ann_type in {ImageAnnotationTypeName.OBJECT_DETECTION}:
train_idx, test_idx = train_test_split(list(range(n)),
train_size=fraction,
random_state=42)
else:
labels = [d['category_id'] for d in lst]
train_idx, test_idx, train_label, test_label = train_test_split(
list(range(n)), labels, stratify=labels, train_size=fraction)
logger.info(f'train idx: {train_idx}.')
logger.info(f'test idx: {test_idx}.')
return train_idx, test_idx
def entrance(mode,
input_transform_path='/mnt/chjinche/test_data/init_transform/',
input_image_path='/mnt/chjinche/test_data/image_dir_test/',
output_path='/mnt/chjinche/test_data/transform_test/'):
params = {'Mode': mode}
ports = {
'Input image transformation':
ImageTransformationDirectory.load(input_transform_path),
'Input image directory':
ImageDirectory.load(input_image_path)
}
task = ApplyImageTransformation()
task.onstart(ports=ports, params=params)
task.apply(ports=ports, params=params).dump(output_path)
logger.info("Transformed dir dumped")
def evaluate(model, loader, print_freq=1, is_test=False):
batch_time = AverageMeter()
losses = AverageMeter()
error = AverageMeter()
# Model on eval mode
model.eval()
end = time.time()
target_list = []
pred_top1_list = []
with torch.no_grad():
for batch_idx, (input, target) in enumerate(loader):
# Create variables
if torch.cuda.is_available():
input = input.cuda()
target = target.cuda()
# compute output
output = model(input)
loss = torch.nn.functional.cross_entropy(output, target)
# measure accuracy and record loss
batch_size = target.size(0)
_, pred = output.data.cpu().topk(1, dim=1)
target_list += target.tolist()
pred_top1_list += [i[0] for i in pred.tolist()]
error.update(
torch.ne(pred.squeeze(), target.cpu()).float().sum().item() /
batch_size, batch_size)
losses.update(loss.item() / batch_size, batch_size)
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
# print stats
if batch_idx % print_freq == 0:
res = '\t'.join([
'Test' if is_test else 'Valid',
'Iter: [{:d}/{:d}]'.format(batch_idx + 1, len(loader)),
'Avg_Time_Batch/Avg_Time_Epoch: {:.3f}/{:.3f}'.format(
batch_time.val, batch_time.avg),
'Avg_Loss_Batch/Avg_Loss_Epoch: {:.4f}/{:.4f}'.format(
losses.val, losses.avg),
'Avg_Error_Batch/Avg_Error_Epoch: {:.4f}/{:.4f}'.format(
error.val, error.avg),
])
logger.info(res)
# Return summary statistics
return batch_time.avg, losses.avg, error.avg, (target_list, pred_top1_list)
def entrance(
mode,
input_transform_path='/mnt/chjinche/test_data/detection/init_transform/',
input_image_path='/mnt/chjinche/test_data/detection/image_dir/',
output_path='/mnt/chjinche/test_data/detection/transform/'):
kwargs = {
'mode':
mode,
'input_image_transformation':
ImageTransformationDirectory.load(input_transform_path),
'input_image_directory':
ImageDirectory.load(input_image_path)
}
task = ApplyImageTransformation()
task.on_init(**kwargs)
output_dir, = task.run(**kwargs)
output_dir.dump(output_path)
logger.info("Transformed dir dumped")
def split_images(src_path, tgt_train_path, tgt_test_path, fraction):
loaded_dir = ImageDirectory.load(src_path)
lst = loaded_dir.image_lst
logger.info(f'Start splitting.')
train_set_lst, test_set_lst = get_stratified_split_list(lst, fraction)
logger.info(f'Got stratified split list. train {len(train_set_lst)}, test {len(test_set_lst)}.')
train_set_dir = FolderBasedImageDirectory.create_with_lst(src_path, train_set_lst)
test_set_dir = FolderBasedImageDirectory.create_with_lst(src_path, test_set_lst)
logger.info('Dump train set.')
train_set_dir.dump(tgt_train_path)
logger.info('Dump test set.')
test_set_dir.dump(tgt_test_path)
def split_images(src_path, tgt_train_path, tgt_test_path, fraction):
# TODO: use multi-label stratified split
# from skmultilearn.model_selection import iterative_train_test_split
# X_train, y_train, X_test, y_test = iterative_train_test_split(X, y, test_size = 0.5)
loaded_dir = ImageDirectory.load(src_path)
ann_type = loaded_dir.get_annotation_type()
logger.info(f'task: {ann_type}')
lst = loaded_dir.image_lst
logger.info(f'Start splitting.')
train_set_idx, test_set_idx = get_split_list(ann_type, lst, fraction)
logger.info(
f'Got split list. train {len(train_set_idx)}, test {len(test_set_idx)}.'
)
train_set_dir = loaded_dir.get_sub_dir(train_set_idx)
test_set_dir = loaded_dir.get_sub_dir(test_set_idx)
logger.info('Dump train set.')
train_set_dir.dump(tgt_train_path)
logger.info('Dump test set.')
test_set_dir.dump(tgt_test_path)
def train_one_epoch(self, loader, optimizer, epoch, epochs, print_freq=1):
batch_time = AverageMeter()
losses = AverageMeter()
error = AverageMeter()
# Model on train mode
# print(self.model)
self.model.train()
end = time.time()
batches = len(loader)
for batch_idx, (input, target) in enumerate(loader):
# Create variables
if torch.cuda.is_available():
input = input.cuda()
target = target.cuda()
# Compute output
output = self.model(input)
loss = torch.nn.functional.cross_entropy(output, target)
# Measure accuracy and record loss
batch_size = target.size(0)
_, pred = output.data.cpu().topk(1, dim=1)
error.update(
torch.ne(pred.squeeze(), target.cpu()).float().sum().item() /
batch_size, batch_size)
losses.update(loss.item(), batch_size)
# Compute gradient and do SGD step
optimizer.zero_grad()
loss.backward()
optimizer.step()
# Measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
if batch_idx % print_freq == 0:
res = '\t'.join([
f'Epoch: [{epoch + 1}/{epochs}]',
f'Iter: [{batch_idx + 1}/{batches}]',
f'Avg_Time_Batch/Avg_Time_Epoch: {batch_time.val:.3f}/{batch_time.avg:.3f}',
f'Avg_Loss_Batch/Avg_Loss_Epoch: {losses.val:.4f}/{losses.avg:.4f}',
f'Avg_Error_Batch/Avg_Error_Epoch: {error.val:.4f}/{error.avg:.4f}'
])
logger.info(res)
# Return summary statistics
return batch_time.avg, losses.avg, error.avg
def split_images(src_path, tgt_train_path, tgt_test_path, fraction):
loaded_dir = ImageDirectory.load(src_path)
lst = loaded_dir.image_lst
logger.info(f'Start splitting.')
train_set_idx, test_set_idx = get_stratified_split_list(lst, fraction)
logger.info(
f'Got stratified split list. train {len(train_set_idx)}, test {len(test_set_idx)}.'
)
train_set_dir = loaded_dir.get_sub_dir(train_set_idx)
test_set_dir = loaded_dir.get_sub_dir(test_set_idx)
logger.info('Dump train set.')
train_set_dir.dump(tgt_train_path)
logger.info('Dump test set.')
test_set_dir.dump(tgt_test_path)
def __init__(self,
model_type='densenet201',
pretrained=True,
memory_efficient=False,
num_classes=20):
logger.info('Init DenseNet.')
super(DenseNet, self).__init__()
if not pretrained:
model_type = "densenet201"
densenet_func = getattr(densenet, model_type, None)
if densenet_func is None:
# todo: catch exception and throw AttributeError
logger.info(f"Error: No such pretrained model {model_type}")
sys.exit()
logger.info(f"Model type {model_type}, pretrained {pretrained}")
self.model1 = densenet_func(pretrained=pretrained)
# Pretrained model is trained based on 1000-class ImageNet dataset
self.model2 = nn.Linear(1000, num_classes)
logger.info(f"Model init finished")
def entrance(save_model_path='/mnt/chjinche/test_data/detection/init_model',
model_type='fasterrcnn_resnet50_fpn',
pretrained=True):
model_config = {
'model_class': 'FasterRCNN',
'model_type': model_type,
'pretrained': pretrained
}
logger.info('Dump untrained model.')
logger.info(f'Model config: {model_config}.')
dumper = pickle_dumper(model_config, 'model_config.pkl')
save_model_to_directory(save_model_path, dumper)
logger.info('Finished.')
def entrance(
save_model_path='/mnt/chjinche/projects/saved_untrained_model_resnet',
model_type='resnext101_32x8d',
pretrained=True):
model_config = {
'model_class': 'ResNet',
'model_type': model_type,
'pretrained': pretrained
}
logger.info('Dump untrained model.')
logger.info(f'Model config: {model_config}.')
dumper = pickle_dumper(model_config, 'model_config.pkl')
save_model_to_directory(save_model_path, dumper)
logger.info('Finished.')
def entrance(save_model_path='../init_model',
model_type='densenet201',
pretrained=True,
memory_efficient=False):
model_config = {
'model_class': 'DenseNet',
'model_type': model_type,
'pretrained': pretrained,
'memory_efficient': memory_efficient
}
logger.info('Dump untrained model.')
logger.info(f'Model config: {model_config}.')
dumper = pickle_dumper(model_config, 'model_config.pkl')
save_model_to_directory(save_model_path, dumper)
logger.info('Finished.')
def entrance(
input_path='/mnt/chjinche/test_data/classification/compressed/image_dataset.zip',
output_path='/mnt/chjinche/test_data/test/image_dir/'):
logger.info('Start!')
logger.info(f'input path {input_path}')
# TODO:Case 1: input path is torchvision ImageFolder
# loader_dir = FolderBasedImageDirectory.load_organized(input_path)
# Case 2: input path is compressed file
compressed_extensions = {'.tar', '.zip'}
compressed_path = None
if Path(input_path).is_file():
compressed_path = input_path
else:
for path in Path(input_path).glob(r'**/*'):
print(path)
if path.suffix in compressed_extensions:
compressed_path = path
logger.info(f'compressed file path {compressed_path}')
loader_dir = ImageDirectory.load_compressed(compressed_path)
# TODO: Case 3: input path is custom directory
loader_dir.dump(output_path)
logger.info('Finished.')
def fit(self,
train_set,
valid_set,
epochs,
batch_size,
lr=0.001,
wd=0.0001,
momentum=0.9,
random_seed=None,
patience=10):
logger.info('Torch cuda random seed setting.')
# Torch cuda random seed setting
if random_seed is not None:
if torch.cuda.is_available():
if torch.cuda.device_count() > 1:
torch.cuda.manual_seed_all(random_seed)
else:
torch.cuda.manual_seed(random_seed)
else:
torch.manual_seed(random_seed)
logger.info("Data start loading.")
# DataLoader
train_loader = torch.utils.data.DataLoader(
train_set,
batch_size=batch_size,
shuffle=True,
pin_memory=(torch.cuda.is_available()),
num_workers=0)
valid_loader = torch.utils.data.DataLoader(
valid_set,
batch_size=batch_size,
shuffle=False,
pin_memory=(torch.cuda.is_available()),
num_workers=0)
if torch.cuda.is_available():
self.model = self.model.cuda()
if torch.cuda.device_count() > 1:
# self.model = torch.nn.parallel.DistributedDataParallel(self.model).cuda()
self.model = torch.nn.DataParallel(self.model).cuda()
optimizer = torch.optim.SGD(self.model.parameters(),
lr=lr,
momentum=momentum,
nesterov=True,
weight_decay=wd)
scheduler = torch.optim.lr_scheduler.MultiStepLR(
optimizer, milestones=[0.5 * epochs, 0.75 * epochs], gamma=0.1)
logger.info('Start training epochs.')
best_error = 1
counter = 0
for epoch in range(epochs):
_, train_loss, train_error = self.train_one_epoch(
loader=train_loader,
optimizer=optimizer,
epoch=epoch,
epochs=epochs)
scheduler.step(epoch=epoch)
_, valid_loss, valid_error, _ = evaluate(model=self.model,
loader=valid_loader)
# Determine if model is the best
if valid_error < best_error:
is_best = True
best_error = valid_error
else:
is_best = False
# Early stop
last_epoch_valid_loss = 0
if epoch == 0:
last_epoch_valid_loss = valid_loss
else:
if valid_loss >= last_epoch_valid_loss:
counter += 1
else:
counter = 0
last_epoch_valid_loss = valid_loss
logger.info(
f'valid loss did not decrease consecutively for {counter} epoch'
)
# TODO: save checkpoint files, but removed now to increase web service deployment efficiency.
logger.info(','.join([
f'Epoch {epoch + 1:d}', f'train_loss {train_loss:.6f}',
f'train_error {train_error:.6f}',
f'valid_loss {valid_loss:.5f}',
f'valid_error {valid_error:.5f}'
]))
if is_best:
logger.info(
# f'Get better top1 accuracy: {1-best_error:.4f} will saving weights to {best_checkpoint_name}'
f'Get better top1 accuracy: {1-best_error:.4f}, best checkpoint will be updated.'
)
early_stop = True if counter >= patience else False
if early_stop:
logger.info("Early stopped.")
break
def __init__(self, **kwargs):
super().__init__()
self.kwargs = kwargs.copy()
self.pretrained = kwargs.get('pretrained', None)
# TODO: error if pretrained is None
self.model_type = kwargs.get('model_type', None)
logger.info(f'Model config {kwargs}.')
kwargs.pop('model_type', None)
logger.info(f'Init {self.model_type}.')
model_func = getattr(models, self.model_type, None)
if model_func is None:
# TODO: catch exception and throw AttributeError
logger.info(f"Error: No such pretrained model {self.model_type}")
sys.exit()
if self.pretrained:
# Drop 'num_classes' para to avoid size mismatch
kwargs.pop('num_classes', None)
logger.info(f'Model config {kwargs}.')
self.model = model_func(**kwargs)
else:
logger.info(f'Model config {kwargs}.')
self.model = model_func(**kwargs)
logger.info(f"Model init finished, {self.model}.")
def entrance(input_model_path='../init_model',
train_data_path='../transform_train/',
valid_data_path='../transform_test/',
save_model_path='../saved_model',
epochs=20,
batch_size=16,
learning_rate=0.001,
random_seed=231,
patience=2):
logger.info("Start training.")
logger.info(f"data path: {train_data_path}")
logger.info(f"data path: {valid_data_path}")
train_set = ImageDirectory.load(train_data_path).to_torchvision_dataset()
logger.info(f"Training classes: {train_set.classes}")
valid_set = ImageDirectory.load(valid_data_path).to_torchvision_dataset()
# TODO: assert the same classes between train_set and valid_set.
logger.info("Made dataset")
classes = train_set.classes
num_classes = len(classes)
# TODO: use image directory api to get id-to-class mapping.
id_to_class_dict = {i: classes[i] for i in range(num_classes)}
logger.info("Start building model.")
model_config = load_model_from_directory(input_model_path,
model_loader=pickle_loader).data
model_class = getattr(modellib, model_config.get('model_class', None),
None)
logger.info(f'Model class: {model_class}.')
model_config.pop('model_class', None)
model_config['num_classes'] = num_classes
logger.info(f'Model_config: {model_config}.')
model = model_class(**model_config)
logger.info("Built model. Start training.")
model = train(model=model,
train_set=train_set,
valid_set=valid_set,
epochs=epochs,
batch_size=batch_size,
lr=learning_rate,
random_seed=random_seed,
patience=patience)
# Save model file, configs and install dependencies
# TODO: designer.model could support pathlib.Path
conda = {
"dependencies": [{
"pip": [
"azureml-defaults",
"azureml-designer-core[image]==0.0.25.post7829218",
"fire==0.1.3",
"git+https://github.com/StudioCommunity/CustomModules-1.git@master#subdirectory=azureml-custom-module-examples/image-classification",
"--extra-index-url=https://azureml-modules:3nvdtawseij7o2oenxojj35c43i5lu2ucf77pugohh4g5eqn6xnq@msdata.pkgs.visualstudio.com/_packaging/azureml-modules%40Local/pypi/simple/"
]
}]
}
logger.info(f"Saving with conda: {conda}")
save_pytorch_state_dict_model(model,
init_params=model_config,
path=save_model_path,
task_type=TaskType.MultiClassification,
label_map=id_to_class_dict,
conda=conda)
logger.info('This experiment has been completed.')
def entrance(train_data_path='/mnt/chjinche/data/out_transform_train/',
valid_data_path='/mnt/chjinche/data/out_transform_test/',
save_model_path='/mnt/chjinche/projects/saved_model',
model_type='densenet201',
pretrained=True,
memory_efficient=False,
epochs=1,
batch_size=16,
learning_rate=0.001,
random_seed=231,
patience=2):
logger.info("Start training.")
logger.info(f"data path: {train_data_path}")
logger.info(f"data path: {valid_data_path}")
train_set = ImageDirectory.load(train_data_path).to_torchvision_dataset()
logger.info(f"Training classes: {train_set.classes}")
valid_set = ImageDirectory.load(valid_data_path).to_torchvision_dataset()
# assert the same classes between train_set and valid_set.
logger.info("Made dataset")
classes = train_set.classes
num_classes = len(classes)
# TODO: use image directory api to get id-to-class mapping.
id_to_class_dict = {i: classes[i] for i in range(num_classes)}
logger.info("Start constructing model")
model_config = {
'model_type': model_type,
'pretrained': pretrained,
'memory_efficient': memory_efficient,
'num_classes': num_classes
}
model = DenseNet(**model_config)
model = train(model=model,
train_set=train_set,
valid_set=valid_set,
epochs=epochs,
batch_size=batch_size,
lr=learning_rate,
random_seed=random_seed,
patience=patience)
# Save model file, configs and install dependencies
# TODO: designer.model could support pathlib.Path
conda = {
"dependencies": [{
"pip": [
"azureml-defaults",
"azureml-designer-core[image]==0.0.25.post7829218",
"fire==0.1.3",
"git+https://github.com/StudioCommunity/CustomModules-1.git@master#subdirectory=azureml-custom-module-examples/image-classification",
"--extra-index-url=https://azureml-modules:3nvdtawseij7o2oenxojj35c43i5lu2ucf77pugohh4g5eqn6xnq@msdata.pkgs.visualstudio.com/_packaging/azureml-modules%40Local/pypi/simple/"
]
}]
}
save_pytorch_state_dict_model(model,
init_params=model_config,
path=save_model_path,
task_type=TaskType.MultiClassification,
label_map=id_to_class_dict,
conda=conda)
logger.info('This experiment has been completed.')
def fit(self,
train_set,
valid_set,
epochs,
batch_size,
lr=0.001,
wd=0.0001,
momentum=0.9,
random_seed=None,
patience=10):
logger.info('Torch cuda random seed setting.')
# Torch cuda random seed setting
if random_seed is not None:
if torch.cuda.is_available():
if torch.cuda.device_count() > 1:
torch.cuda.manual_seed_all(random_seed)
else:
torch.cuda.manual_seed(random_seed)
else:
torch.manual_seed(random_seed)
logger.info("Data start loading.")
# DataLoader
# train_set = remove_images_without_annotations(train_set)
train_loader = torch.utils.data.DataLoader(
train_set,
batch_size=batch_size,
shuffle=True,
collate_fn=utils.collate_fn,
pin_memory=(torch.cuda.is_available()),
num_workers=0)
valid_loader = torch.utils.data.DataLoader(
valid_set,
batch_size=batch_size,
shuffle=False,
collate_fn=utils.collate_fn,
pin_memory=(torch.cuda.is_available()),
num_workers=0)
if torch.cuda.is_available():
self.model = self.model.cuda()
if torch.cuda.device_count() > 1:
# self.model = torch.nn.parallel.DistributedDataParallel(self.model).cuda()
self.model = torch.nn.DataParallel(self.model).cuda()
# reduce learning rate every step_size epochs by a factor of gamma (by default) 0.1.
step_size = None
if step_size is None:
step_size = int(np.round(epochs / 1.5))
self.device = torch_device()
# construct our optimizer
params = [p for p in self.model.parameters() if p.requires_grad]
self.optimizer = torch.optim.SGD(
params, lr=lr, momentum=momentum, weight_decay=wd
)
# and a learning rate scheduler
self.lr_scheduler = torch.optim.lr_scheduler.StepLR(
self.optimizer, step_size=step_size, gamma=0.1
)
# store data in these arrays to plot later
self.losses = []
self.ap = []
self.ap_iou_point_5 = []
# main training loop
self.epochs = epochs
for epoch in range(self.epochs):
# train for one epoch, printing every 10 iterations
logger_per_epoch = train_one_epoch(
self.model,
self.optimizer,
train_loader,
self.device,
epoch,
print_freq=10,
)
self.losses.append(logger_per_epoch.meters["loss"].median)
# update the learning rate
self.lr_scheduler.step(epoch=epoch)
# evaluate
e = evaluate(self.model, valid_loader, self.device)
self.ap.append(_calculate_ap(e))
self.ap_iou_point_5.append(_calculate_ap(e, iou_threshold_idx=0))
def __init__(self, **kwargs):
super().__init__(**kwargs)
self.update_nn()
logger.info(f"Model init finished, {self.model}.")
