def test_nested_model_sessions():
model = torch.nn.Sequential(torch.nn.Linear(1, 1), torch.nn.Dropout(0.1))
print(model.training)
with TorchEvalSession(model):
print(model.training)
with TorchTrainSession(model):
print(model.training)
with TorchEvalSession(model):
print(model.training)
with TorchTrainSession(model):
print(model.training)
with TorchEvalSession(model):
print(model.training)
print(model.training)
def run_seg_traced_webcam_demo():
"""
:return:
:rtype:"""
import torch
import io
load_path = (PROJECT_APP_PATH.user_data / "penn_fudan_segmentation" /
"seg_skip_fis").with_suffix(".traced")
# print(load_path)
# torch.jit.load(str(load_path))
with open(str(load_path),
"rb") as f: # Load ScriptModule from io.BytesIO object
buffer = io.BytesIO(f.read())
model = torch.jit.load(buffer) # Load all tensors to the original device
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
])
from matplotlib.pyplot import imshow
with TorchDeviceSession(device=global_torch_device("cpu"), model=model):
with TorchEvalSession(model):
for image in tqdm(frame_generator(cv2.VideoCapture(0))):
result = model(
transform(image).unsqueeze(0).to(global_torch_device()))[0]
imshow(result[0][0].numpy(), vmin=0.0, vmax=1.0)
show()
def maskrcnn_evaluate(
model: Module,
data_loader: DataLoader,
*,
device=global_torch_device(),
writer: Writer = None,
) -> CocoEvaluator:
"""
Args:
model:
data_loader:
device:
writer:
Returns:
"""
n_threads = torch.get_num_threads()
# FIXME remove this and make paste_masks_in_image run on the GPU
torch.set_num_threads(1)
cpu_device = torch.device("cpu")
coco_evaluator = CocoEvaluator(
get_coco_api_from_dataset(data_loader.dataset), get_iou_types(model))
with torch.no_grad():
with TorchEvalSession(model):
for image, targets in tqdm.tqdm(data_loader):
image = [img.to(device) for img in image]
targets = [{k: v.to(device)
for k, v in t.items()} for t in targets]
torch.cuda.synchronize(device)
model_time = time.time()
outputs = model(image)
outputs = [{k: v.to(cpu_device)
for k, v in t.items()} for t in outputs]
model_time = time.time() - model_time
res = {
target["image_id"].item(): output
for target, output in zip(targets, outputs)
}
evaluator_time = time.time()
coco_evaluator.update(res)
evaluator_time = time.time() - evaluator_time
if writer:
writer.scalar("model_time", model_time)
writer.scalar("evaluator_time", evaluator_time)
coco_evaluator.synchronize_between_processes()
coco_evaluator.accumulate()
coco_evaluator.summarize()
torch.set_num_threads(n_threads)
return coco_evaluator
def main(model_name: str = "maskrcnn_pennfudanped", score_threshold=0.55):
base_path = PROJECT_APP_PATH.user_data / 'maskrcnn'
dataset_root = Path.home() / "Data"
torch_seed(3825)
dataset = PennFudanDataset(dataset_root / "PennFudanPed",
Split.Training)
categories = dataset.categories
if True:
model = load_model(model_name=model_name,
model_directory=base_path / 'models')
else:
model = get_pretrained_instance_segmentation_maskrcnn(
dataset.response_channels)
model.to(global_torch_device())
cpu_device = torch.device("cpu")
with torch.no_grad():
with TorchEvalSession(model):
for image in tqdm(
to_tensor_generator(
frame_generator(cv2.VideoCapture(0)),
device=global_torch_device(),
)):
prediction = model(
# torch_vision_normalize_batch_nchw(
uint_hwc_to_chw_float_tensor(image).unsqueeze(0)
# )
)[0]
(boxes, labels, scores) = (
prediction["boxes"].to(cpu_device).numpy(),
prediction["labels"].to(cpu_device).numpy(),
torch.sigmoid(
prediction["scores"]).to(cpu_device).numpy(),
)
indices = scores > score_threshold
cv2.namedWindow(model_name, cv2.WINDOW_NORMAL)
cv2.imshow(
model_name,
draw_bounding_boxes(
quick_to_pil_image(image),
boxes[indices],
labels=labels[indices],
scores=scores[indices],
categories=categories,
))
if cv2.waitKey(1) == 27:
break # esc to quit
def test_model(
model: VAE,
epoch_i: int,
metric_writer: Writer,
loader: DataLoader,
save_images: bool = True,
):
global LOWEST_L
with TorchEvalSession(model):
test_accum_loss = 0
with torch.no_grad():
for i, (original, labels, *_) in enumerate(loader):
original = original.to(global_torch_device())
reconstruction, mean, log_var = model(original)
loss = loss_function(reconstruction, original, mean,
log_var).item()
test_accum_loss += loss
metric_writer.scalar("test_loss", test_accum_loss)
if save_images:
if i == 0:
n = min(original.size(0), 8)
comparison = torch.cat(
[original[:n], reconstruction[:n]])
save_image(
comparison.cpu(), # Torch save images
str(BASE_PATH /
f"reconstruction_{str(epoch_i)}.png"),
nrow=n,
)
"""
scatter_plot_encoding_space(str(BASE_PATH /
f'encoding_space_{str(epoch_i)}.png'),
mean.to('cpu').numpy(),
log_var.to('cpu').numpy(),
labels)
"""
break
# test_loss /= len(loader.dataset)
test_accum_loss /= loader.batch_size
print(f"====> Test set loss: {test_accum_loss:.4f}")
torch.save(model.state_dict(),
BASE_PATH / f"model_state_dict{str(epoch_i)}.pth")
if LOWEST_L > test_accum_loss:
LOWEST_L = test_accum_loss
torch.save(model.state_dict(), BASE_PATH / f"best_state_dict.pth")
def run_seg_traced_webcam_demo():
"""
:param categories:
:type categories:
:param cfg:
:type cfg:
:param model_ckpt:
:type model_ckpt:
:param score_threshold:
:type score_threshold:
:param window_name:
:type window_name:
:return:
:rtype:
"""
import torch
import io
torch.jit.load("seg_skip_fis.traced")
with open(
"seg_skip_fis.traced", "rb"
) as f: # Load ScriptModule from io.BytesIO object
buffer = io.BytesIO(f.read())
model = torch.jit.load(buffer) # Load all tensors to the original device
transform = transforms.Compose(
[
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
]
)
with TorchDeviceSession(
device=global_torch_device(cuda_if_available=False), model=model
):
with TorchEvalSession(model):
for image in tqdm(frame_generator(cv2.VideoCapture(0))):
result = model(transform(image).unsqueeze(0).to(global_torch_device()))[
0
]
print(result)
from matplotlib.pyplot import imshow
imshow(result[0][0].numpy(), vmin=0.0, vmax=1.0)
show()
def validate_model(model, valid_loader):
with TorchDeviceSession(global_torch_device("cpu"), model):
with torch.no_grad():
with TorchCacheSession():
with TorchEvalSession(model):
valid_masks = []
out_data = []
a = (256, 256)
tr = min(len(valid_loader.dataset) * 4, 2000)
probabilities = []
for sample_i, (data,
target) in enumerate(tqdm(valid_loader)):
data = data.to(global_torch_device())
outpu, *_ = model(data)
for m, d, p in zip(
target.cpu().detach().numpy(),
data.cpu().detach().numpy(),
torch.sigmoid(outpu).cpu().detach().numpy(),
):
out_data.append(cv2_resize(chw_to_hwc(d), a))
valid_masks.append(cv2_resize(m[0], a))
probabilities.append(cv2_resize(p[0], a))
sample_i += 1
if sample_i >= tr - 1:
break
if sample_i >= tr - 1:
break
min_a = min(3, len(out_data))
f, ax = pyplot.subplots(min_a, 3, figsize=(24, 12))
# assert len(valid_masks)>2, f'{len(valid_masks), tr}'
for i in range(min_a):
ax[0, i].imshow(out_data[i], vmin=0, vmax=1)
ax[0, i].set_title("Original", fontsize=14)
ax[1, i].imshow(valid_masks[i], vmin=0, vmax=1)
ax[1, i].set_title("Target", fontsize=14)
ax[2, i].imshow(probabilities[i], vmin=0, vmax=1)
ax[2, i].set_title("Prediction", fontsize=14)
pyplot.show()
def single_epoch_evaluation(
model: Module,
evaluation_loader: DataLoader,
subset: Split,
*,
epoch: int = None,
writer: Writer = None,
device: torch.device = global_torch_device()) -> float:
correct = 0
num_batches = len(evaluation_loader)
with TorchEvalSession(model):
for data, target in tqdm(evaluation_loader,
desc=f'{subset} batch #',
total=num_batches):
correct += model(data.to(device)).argmax(dim=-1).squeeze().eq(
target.to(device)).sum().item()
acc = correct / len(evaluation_loader.dataset)
if writer:
writer.scalar(f'{subset}_accuracy', acc, epoch)
return acc
def inference(
model: Module, data_iterator: Iterator, denoise: bool = True, num: int = 3
) -> None:
"""
:param model:
:type model:
:param data_iterator:
:type data_iterator:"""
with torch.no_grad():
with TorchEvalSession(model):
img, target = next(data_iterator)
if denoise:
model_input = img + torch.normal(
mean=0.0, std=0.1, size=img.shape, device=global_torch_device()
)
else:
model_input = img
model_input = torch.clamp(model_input, 0.0, 1.0)
pred, *_ = model(model_input)
plot_side_by_side(
[
pred.squeeze(1)[:num].cpu().numpy(),
model_input.squeeze(1)[:num].cpu().numpy(),
]
)
pyplot.show()
return
for i, (s, j, label) in enumerate(
zip(pred.cpu().numpy(), model_input.cpu().numpy(), target)
):
pyplot.imshow(j[0])
pyplot.title(f"sample_{i}, category: {label}")
pyplot.show()
# plot_side_by_side(s)
pyplot.imshow(s[0])
pyplot.title(f"sample_{i}, category: {label}")
pyplot.show()
break
def train_siamese(
model,
optimiser,
criterion,
*,
writer: Writer = MockWriter(),
train_number_epochs,
data_dir,
train_batch_size,
model_name,
save_path,
save_best=False,
img_size,
validation_interval: int = 1,
):
"""
:param data_dir:
:type data_dir:
:param optimiser:
:type optimiser:
:param criterion:
:type criterion:
:param writer:
:type writer:
:param model_name:
:type model_name:
:param save_path:
:type save_path:
:param save_best:
:type save_best:
:param model:
:type model:
:param train_number_epochs:
:type train_number_epochs:
:param train_batch_size:
:type train_batch_size:
:return:
:rtype:
Parameters
----------
img_size
validation_interval"""
train_dataloader = DataLoader(
TripletDataset(
data_path=data_dir,
transform=transforms.Compose([
transforms.Grayscale(),
transforms.Resize(img_size),
transforms.ToTensor(),
]),
split=SplitEnum.training,
),
shuffle=True,
num_workers=0,
batch_size=train_batch_size,
)
valid_dataloader = DataLoader(
TripletDataset(
data_path=data_dir,
transform=transforms.Compose([
transforms.Grayscale(),
transforms.Resize(img_size),
transforms.ToTensor(),
]),
split=SplitEnum.validation,
),
shuffle=True,
num_workers=0,
batch_size=train_batch_size,
)
best = math.inf
E = tqdm(range(0, train_number_epochs))
batch_counter = count()
for epoch in E:
for tss in train_dataloader:
batch_i = next(batch_counter)
with TorchTrainSession(model):
optimiser.zero_grad()
loss_contrastive = criterion(*model(
*[t.to(global_torch_device()) for t in tss]))
loss_contrastive.backward()
optimiser.step()
a = loss_contrastive.cpu().item()
writer.scalar("train_loss", a, batch_i)
if batch_counter.__next__() % validation_interval == 0:
with TorchEvalSession(model):
for tsv in valid_dataloader:
o = model(*[t.to(global_torch_device()) for t in tsv])
a_v = criterion(*o).cpu().item()
valid_positive_acc = (accuracy(
distances=pairwise_distance(o[0], o[1]),
is_diff=0).cpu().item())
valid_negative_acc = (accuracy(
distances=pairwise_distance(o[0], o[2]),
is_diff=1).cpu().item())
valid_acc = numpy.mean(
(valid_negative_acc, valid_positive_acc))
writer.scalar("valid_loss", a_v, batch_i)
writer.scalar("valid_positive_acc", valid_positive_acc,
batch_i)
writer.scalar("valid_negative_acc", valid_negative_acc,
batch_i)
writer.scalar("valid_acc", valid_acc, batch_i)
if a_v < best:
best = a_v
print(f"new best {best}")
if save_best:
save_model_parameters(
model,
optimiser=optimiser,
model_name=model_name,
save_directory=save_path,
)
E.set_description(
f"Epoch number {epoch}, Current train loss {a}, valid loss {a_v}, valid acc {valid_acc}"
)
return model
def pred_target_train_model(
model,
train_iterator,
criterion,
optimizer,
scheduler,
writer,
interrupted_path,
test_data_iterator=None,
num_updates: int = 250000,
early_stop=None,
) -> torch.nn.Module:
best_model_wts = copy.deepcopy(model.state_dict())
best_val_loss = 1e10
since = time.time()
try:
sess = tqdm.tqdm(range(num_updates), leave=False, disable=False)
val_loss = 0
update_loss = 0
val_acc = 0
last_val = None
last_out = None
with torch.autograd.detect_anomaly():
for update_i in sess:
for phase in [Split.Training, Split.Validation]:
if phase == Split.Training:
with TorchTrainSession(model):
input, true_label = zip(*next(train_iterator))
rgb_imgs = torch_vision_normalize_batch_nchw(
uint_nhwc_to_nchw_float_batch(
rgb_drop_alpha_batch_nhwc(
to_tensor(input))))
true_label = to_tensor(true_label,
dtype=torch.long)
optimizer.zero_grad()
pred = model(rgb_imgs)
loss = criterion(pred, true_label)
loss.backward()
optimizer.step()
if last_out is None:
last_out = pred
else:
if not torch.dist(last_out, pred) > 0:
print(f"Same output{last_out},{pred}")
last_out = pred
update_loss = loss.data.cpu().numpy()
writer.scalar(f"loss/train", update_loss, update_i)
if scheduler:
scheduler.step()
elif test_data_iterator:
with TorchEvalSession(model):
test_rgb_imgs, test_true_label = zip(
*next(train_iterator))
test_rgb_imgs = torch_vision_normalize_batch_nchw(
uint_nhwc_to_nchw_float_batch(
rgb_drop_alpha_batch_nhwc(
to_tensor(test_rgb_imgs))))
test_true_label = to_tensor(test_true_label,
dtype=torch.long)
with torch.no_grad():
val_pred = model(test_rgb_imgs)
val_loss = criterion(val_pred, test_true_label)
_, cat = torch.max(val_pred, -1)
val_acc = torch.sum(
cat == test_true_label) / float(cat.size(0))
writer.scalar(f"loss/acc", val_acc, update_i)
writer.scalar(f"loss/val", val_loss, update_i)
if last_val is None:
last_val = cat
else:
if all(last_val == cat):
print(f"Same val{last_val},{cat}")
last_val = cat
if val_loss < best_val_loss:
best_val_loss = val_loss
best_model_wts = copy.deepcopy(
model.state_dict())
sess.write(
f"New best validation model at update {update_i} with test_loss {best_val_loss}"
)
torch.save(model.state_dict(),
interrupted_path)
if early_stop is not None and val_pred < early_stop:
break
sess.set_description_str(f"Update {update_i} - {phase} "
f"update_loss:{update_loss:2f} "
f"test_loss:{val_loss}"
f"val_acc:{val_acc}")
except KeyboardInterrupt:
print("Interrupt")
finally:
pass
model.load_state_dict(best_model_wts) # load best model weights
time_elapsed = time.time() - since
print(f"{time_elapsed // 60:.0f}m {time_elapsed % 60:.0f}s")
print(f"Best val loss: {best_val_loss:3f}")
return model
def run_webcam_demo(
cfg: NOD,
input_cfg: NOD,
categories: List,
model_ckpt: Path,
score_threshold: float = 0.7,
window_name: str = "SSD",
):
"""
:param categories:
:type categories:
:param cfg:
:type cfg:
:param model_ckpt:
:type model_ckpt:
:param score_threshold:
:type score_threshold:
:param window_name:
:type window_name:
:return:
:rtype:
"""
cpu_device = torch.device("cpu")
transforms = SSDTransform(input_cfg.image_size,
input_cfg.pixel_mean,
split=Split.Testing)
model = SingleShotDectectionNms(cfg)
checkpointer = CheckPointer(model,
save_dir=ensure_existence(
PROJECT_APP_PATH.user_data / "results"))
checkpointer.load(model_ckpt, use_latest=model_ckpt is None)
print(
f"Loaded weights from {model_ckpt if model_ckpt else checkpointer.get_checkpoint_file()}"
)
model.post_init()
model.to(global_torch_device())
with TorchEvalSession(model):
for image in tqdm(frame_generator(cv2.VideoCapture(0))):
result = model(
transforms(image)[0].unsqueeze(0).to(global_torch_device()))
height, width, *_ = image.shape
result.boxes[:, 0::2] *= width / result.img_width.cpu().item()
result.boxes[:, 1::2] *= height / result.img_height.cpu().item()
(boxes, labels, scores) = (
result.boxes.to(cpu_device).numpy(),
result.labels.to(cpu_device).numpy(),
result.scores.to(cpu_device).numpy(),
)
indices = scores > score_threshold
cv2.namedWindow(window_name, cv2.WINDOW_NORMAL)
cv2.imshow(
window_name,
draw_bounding_boxes(
image,
boxes[indices],
labels=labels[indices],
scores=scores[indices],
categories=categories,
score_font=ImageFont.truetype(
PACKAGE_DATA_PATH / "Lato-Regular.ttf",
24,
),
).astype(numpy.uint8),
)
if cv2.waitKey(1) == 27:
break # esc to quit
def predictor_response_train_model(
model,
*,
train_iterator,
criterion,
optimizer,
scheduler,
writer,
interrupted_path,
val_data_iterator=None,
num_updates: int = 250000,
device=global_torch_device(),
early_stop=None,
):
"""
:param model:
:param train_iterator:
:param criterion:
:param optimizer:
:param scheduler:
:param writer:
:param interrupted_path:
:param val_data_iterator:
:param num_updates:
:param device:
:param early_stop:
:return:
"""
best_model_wts = copy.deepcopy(model.state_dict())
best_val_loss = 1e10
since = time.time()
try:
sess = tqdm(range(num_updates), leave=False, disable=False)
val_loss = 0
update_loss = 0
val_acc = 0
last_val = None
last_out = None
with torch.autograd.detect_anomaly():
for update_i in sess:
for phase in [Split.Training, Split.Validation]:
if phase == Split.Training:
with TorchTrainSession(model):
input, true_label = next(train_iterator)
rgb_imgs = to_tensor(
input, dtype=torch.float, device=device
).repeat(1, 3, 1, 1)
true_label = to_tensor(
true_label, dtype=torch.long, device=device
)
optimizer.zero_grad()
pred = model(rgb_imgs)
loss = criterion(pred, true_label)
loss.backward()
optimizer.step()
update_loss = loss.data.cpu().numpy()
writer.scalar(f"loss/train", update_loss, update_i)
if scheduler:
scheduler.step()
elif val_data_iterator:
with TorchEvalSession(model):
test_rgb_imgs, test_true_label = next(val_data_iterator)
test_rgb_imgs = to_tensor(
test_rgb_imgs, dtype=torch.float, device=device
).repeat(1, 3, 1, 1)
test_true_label = to_tensor(
test_true_label, dtype=torch.long, device=device
)
with torch.no_grad():
val_pred = model(test_rgb_imgs)
val_loss = criterion(val_pred, test_true_label)
_, cat = torch.max(val_pred, -1)
val_acc = torch.sum(cat == test_true_label) / float(
cat.size(0)
)
writer.scalar(f"loss/acc", val_acc, update_i)
writer.scalar(f"loss/val", val_loss, update_i)
if val_loss < best_val_loss:
best_val_loss = val_loss
best_model_wts = copy.deepcopy(model.state_dict())
sess.write(
f"New best validation model at update {update_i} with best_val_loss {best_val_loss}"
)
torch.save(model.state_dict(), interrupted_path)
if early_stop is not None and val_pred < early_stop:
break
sess.set_description_str(
f"Update {update_i} - {phase} "
f"update_loss:{update_loss:2f} "
f"val_loss:{val_loss}"
f"val_acc:{val_acc}"
)
except KeyboardInterrupt:
print("Interrupt")
finally:
pass
model.load_state_dict(best_model_wts) # load best model weights
time_elapsed = time.time() - since
print(f"{time_elapsed // 60:.0f}m {time_elapsed % 60:.0f}s")
print(f"Best val loss: {best_val_loss:3f}")
return model
def train_person_segmentor(
model: torch.nn.Module,
train_loader: torch.utils.data.DataLoader,
valid_loader: torch.utils.data.DataLoader,
criterion: callable,
optimiser: torch.optim.Optimizer,
*,
save_model_path: Path,
learning_rate: Number = 6e-2,
scheduler: torch.optim.lr_scheduler = None,
n_epochs: int = 100,
writer: ImageWriterMixin = MockWriter(),
):
"""
:param model:
:type model:
:param train_loader:
:type train_loader:
:param valid_loader:
:type valid_loader:
:param criterion:
:type criterion:
:param optimiser:
:type optimiser:
:param scheduler:
:type scheduler:
:param save_model_path:
:type save_model_path:
:param n_epochs:
:type n_epochs:
:return:
:rtype:"""
valid_loss_min = numpy.Inf # track change in validation loss
assert n_epochs > 0, n_epochs
E = tqdm(range(1, n_epochs + 1))
for epoch_i in E:
train_loss = 0.0
valid_loss = 0.0
with TorchTrainSession(model):
for data, target in tqdm(train_loader):
output, *_ = model(data.to(global_torch_device()))
loss = criterion(output,
target.to(global_torch_device()).float())
optimiser.zero_grad()
loss.backward()
optimiser.step()
train_loss += loss.cpu().item() * data.size(0)
with TorchEvalSession(model):
with torch.no_grad():
for data, target in tqdm(valid_loader):
target = target.float()
(
output,
*_,
) = model( # forward pass: compute predicted outputs by passing inputs to the model
data.to(global_torch_device()))
validation_loss = criterion(
output, target.to(
global_torch_device())) # calculate the batch loss
writer.scalar(
"dice_validation",
dice_loss(output, target.to(global_torch_device())),
)
valid_loss += validation_loss.detach().cpu().item(
) * data.size(0) # update average validation loss
writer.image("input", data, epoch_i) # write the last batch
writer.image("truth", target, epoch_i) # write the last batch
writer.image("prediction", torch.sigmoid(output),
epoch_i) # write the last batch
# calculate average losses
train_loss = train_loss / len(train_loader.dataset)
valid_loss = valid_loss / len(valid_loader.dataset)
# save model if validation loss has decreased
if valid_loss <= valid_loss_min:
print(
f"Validation loss decreased ({valid_loss_min:.6f} --> {valid_loss:.6f}). Saving model ..."
)
torch.save(model.state_dict(), save_model_path)
valid_loss_min = valid_loss
if scheduler:
scheduler.step()
optimiser, scheduler = reschedule_learning_rate(
model,
optimiser,
epoch_i,
scheduler,
starting_learning_rate=learning_rate,
)
# print training/validation statistics
current_lr = next(iter(optimiser.param_groups))["lr"]
E.set_description(f"Epoch: {epoch_i} "
f"Training Loss: {train_loss:.6f} "
f"Validation Loss: {valid_loss:.6f} "
f"Learning rate: {current_lr:.6f}")
writer.scalar("training_loss", train_loss)
writer.scalar("validation_loss", valid_loss)
writer.scalar("learning_rate", current_lr)
return model
def inner_train_ssd(*,
data_root: Path,
cfg: NOD,
model: Module,
data_loader: DataLoader,
optimiser: Optimizer,
scheduler: WarmupMultiStepLR,
check_pointer: callable,
device: callable,
arguments: callable,
kws: NOD,
) -> Module:
"""
:param data_root:
:type data_root:
:param cfg:
:type cfg:
:param model:
:type model:
:param data_loader:
:type data_loader:
:param optimiser:
:type optimiser:
:param scheduler:
:type scheduler:
:param check_pointer:
:type check_pointer:
:param device:
:type device:
:param arguments:
:type arguments:
:param kws:
:type kws:
:return:
:rtype:
"""
logger = logging.getLogger("SSD.trainer")
logger.info("Start training ...")
meters = MetricLogger()
with TorchTrainSession(model):
save_to_disk = global_distribution_rank() == 0
if kws.use_tensorboard and save_to_disk:
import tensorboardX
writer = tensorboardX.SummaryWriter(
log_dir=str(PROJECT_APP_PATH.user_data / "results" / "tf_logs")
)
else:
writer = None
max_iter = len(data_loader)
start_iter = arguments["iteration"]
start_training_time = time.time()
end = time.time()
for iteration, (images, targets, _) in enumerate(data_loader, start_iter):
arguments["iteration"] = iteration
images = images.to(device)
targets = targets.to(device)
loss_instance = MultiBoxLoss(neg_pos_ratio=cfg.model.neg_pos_ratio)
cls_logits, bbox_pred = model(images)
reg_loss, cls_loss = loss_instance(
cls_logits, bbox_pred, targets.labels, targets.boxes
)
loss_dict = dict(reg_loss=reg_loss, cls_loss=cls_loss)
loss = sum(loss for loss in loss_dict.values())
loss_dict_reduced = reduce_loss_dict(
loss_dict
) # reduce losses over all GPUs for logging purposes
losses_reduced = sum(loss for loss in loss_dict_reduced.values())
meters.update(total_loss=losses_reduced, **loss_dict_reduced)
optimiser.zero_grad()
loss.backward()
optimiser.step()
scheduler.step()
batch_time = time.time() - end
end = time.time()
meters.update(time=batch_time)
if iteration % kws.log_step == 0:
eta_seconds = meters.time.global_avg * (max_iter - iteration)
eta_string = str(datetime.timedelta(seconds=int(eta_seconds)))
logger.info(
meters.delimiter.join(
[
f"iter: {iteration:06d}",
f"lr: {optimiser.param_groups[0]['lr']:.5f}",
f"{str(meters)}",
f"eta: {eta_string}",
f"mem: {round(torch.cuda.max_memory_allocated() / 1024.0 / 1024.0)}M",
]
)
)
if writer:
global_step = iteration
writer.add_scalar(
"losses/total_loss", losses_reduced, global_step=global_step
)
for loss_name, loss_item in loss_dict_reduced.items():
writer.add_scalar(
f"losses/{loss_name}", loss_item, global_step=global_step
)
writer.add_scalar(
"lr", optimiser.param_groups[0]["lr"], global_step=global_step
)
if iteration % kws.save_step == 0:
check_pointer.save(f"model_{iteration:06d}", **arguments)
if (
kws.eval_step > 0
and iteration % kws.eval_step == 0
and not iteration == max_iter
):
with TorchEvalSession(model):
eval_results = do_ssd_evaluation(
data_root,
cfg,
model,
distributed=kws.distributed,
iteration=iteration,
)
if global_distribution_rank() == 0 and writer:
for eval_result, dataset in zip(
eval_results, cfg.datasets.test
):
write_metrics_recursive(
eval_result["metrics"],
"metrics/" + dataset,
writer,
iteration,
)
check_pointer.save("model_final", **arguments)
total_training_time = int(
time.time() - start_training_time
) # compute training time
logger.info(
f"Total training time: {datetime.timedelta(seconds=total_training_time)} ("
f"{total_training_time / max_iter:.4f} s / it)"
)
return model
def main():
""" """
args = argparse.ArgumentParser()
args.add_argument("--inference", "-i", action="store_true")
args.add_argument("--continue_training", "-c", action="store_true")
args.add_argument("--no_cuda", "-k", action="store_false")
args.add_argument("--export", "-e", action="store_true")
options = args.parse_args()
timeas = str(time.time())
this_model_path = PROJECT_APP_PATH.user_data / timeas
this_log = PROJECT_APP_PATH.user_log / timeas
ensure_directory_exist(this_model_path)
ensure_directory_exist(this_log)
best_model_name = "best_validation_model.model"
interrupted_path = str(this_model_path / best_model_name)
torch.manual_seed(seed)
if not options.no_cuda:
global_torch_device("cpu")
with MixedObservationWrapper() as env:
env.seed(seed)
train_iter = batch_generator(iter(env), batch_size)
num_categories = env.sensor("Class").space.discrete_steps
test_iter = train_iter
model, params_to_update = squeezenet_retrain(num_categories)
print(params_to_update)
model = model.to(global_torch_device())
if options.continue_training:
_list_of_files = list(PROJECT_APP_PATH.user_data.rglob("*.model"))
latest_model_path = str(max(_list_of_files, key=os.path.getctime))
print(f"loading previous model: {latest_model_path}")
if latest_model_path is not None:
model.load_state_dict(torch.load(latest_model_path))
criterion = torch.nn.CrossEntropyLoss().to(global_torch_device())
optimiser_ft = optim.SGD(model.parameters(),
lr=learning_rate,
momentum=momentum,
weight_decay=wd)
exp_lr_scheduler = torch.optim.lr_scheduler.StepLR(optimiser_ft,
step_size=7,
gamma=0.1)
with TensorBoardPytorchWriter(this_log) as writer:
if True:
model = pred_target_train_model(
model,
train_iter,
criterion,
optimiser_ft,
exp_lr_scheduler,
writer,
interrupted_path,
test_data_iterator=test_iter,
num_updates=num_updates,
)
inputs, true_label = zip(*next(train_iter))
rgb_imgs = torch_vision_normalize_batch_nchw(
uint_nhwc_to_nchw_float_batch(
rgb_drop_alpha_batch_nhwc(to_tensor(inputs))))
predicted = torch.argmax(model(rgb_imgs), -1)
true_label = to_tensor(true_label, dtype=torch.long)
print(predicted, true_label)
horizontal_imshow(
inputs,
[
f"p:{int(p)},t:{int(t)}"
for p, t in zip(predicted, true_label)
],
)
pyplot.show()
torch.cuda.empty_cache()
if options.export:
with TorchEvalSession(model):
example = torch.rand(1, 3, 256, 256)
traced_script_module = torch.jit.trace(model.to("cpu"), example)
traced_script_module.save("resnet18_v.model")
def train_person_segmenter(
model,
train_loader,
valid_loader,
criterion,
optimizer,
scheduler,
save_model_path: Path,
n_epochs: int = 100,
):
"""
:param model:
:type model:
:param train_loader:
:type train_loader:
:param valid_loader:
:type valid_loader:
:param criterion:
:type criterion:
:param optimizer:
:type optimizer:
:param scheduler:
:type scheduler:
:param save_model_path:
:type save_model_path:
:param n_epochs:
:type n_epochs:
:return:
:rtype:
"""
valid_loss_min = numpy.Inf # track change in validation loss
assert n_epochs > 0, n_epochs
E = tqdm(range(1, n_epochs + 1))
for epoch in E:
train_loss = 0.0
valid_loss = 0.0
dice_score = 0.0
with TorchTrainSession(model):
for data, target in tqdm(train_loader):
data, target = (
data.to(global_torch_device()),
target.to(global_torch_device()),
)
optimizer.zero_grad()
output, *_ = model(data)
output = torch.sigmoid(output)
loss = criterion(output, target)
loss.backward()
optimizer.step()
train_loss += loss.item() * data.size(0)
with TorchEvalSession(model):
with torch.no_grad():
for data, target in tqdm(valid_loader):
data, target = (
data.to(global_torch_device()),
target.to(global_torch_device()),
)
output, *_ = model(
data
) # forward pass: compute predicted outputs by passing inputs to the model
output = torch.sigmoid(output)
loss = criterion(output,
target) # calculate the batch loss
valid_loss += loss.item() * data.size(
0) # update average validation loss
dice_cof = intersection_over_union(
output.cpu().detach().numpy(),
target.cpu().detach().numpy())
dice_score += dice_cof * data.size(0)
# calculate average losses
train_loss = train_loss / len(train_loader.dataset)
valid_loss = valid_loss / len(valid_loader.dataset)
dice_score = dice_score / len(valid_loader.dataset)
# print training/validation statistics
E.set_description(f"Epoch: {epoch}"
f" Training Loss: {train_loss:.6f} "
f"Validation Loss: {valid_loss:.6f} "
f"Dice Score: {dice_score:.6f}")
# save model if validation loss has decreased
if valid_loss <= valid_loss_min:
print(
f"Validation loss decreased ({valid_loss_min:.6f} --> {valid_loss:.6f}). Saving model ..."
)
torch.save(model.state_dict(), save_model_path)
valid_loss_min = valid_loss
scheduler.step()
model, scheduler = reschedule_learning_rate(model, epoch, scheduler)
return model
def main():
pyplot.style.use("bmh")
base_path = Path.home() / "/Data" / "PennFudanPed"
save_model_path = PROJECT_APP_PATH.user_data / 'models' / "penn_fudan_ped_seg.model"
train_model = False
eval_model = not train_model
SEED = 87539842
batch_size = 8
num_workers = 1 # os.cpu_count()
learning_rate = 0.01
torch_seed(SEED)
train_set = PennFudanDataset(base_path, Split.Training)
train_loader = DataLoader(train_set,
batch_size=batch_size,
shuffle=True,
num_workers=num_workers)
valid_loader = DataLoader(
PennFudanDataset(base_path, Split.Validation),
batch_size=batch_size,
shuffle=False,
num_workers=num_workers,
)
model = SkipHourglassFission(
input_channels=train_set.predictor_shape[-1],
output_heads=(train_set.response_shape[-1], ),
encoding_depth=1,
)
model.to(global_torch_device())
if train_model:
if save_model_path.exists():
model.load_state_dict(torch.load(str(save_model_path)))
print("loading saved model")
with TorchTrainSession(model):
criterion = BCEDiceLoss(eps=1.0)
optimiser = torch.optim.SGD(model.parameters(), lr=learning_rate)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
optimiser, T_max=7, eta_min=learning_rate / 100, last_epoch=-1)
model = train_person_segmenter(
model,
train_loader,
valid_loader,
criterion,
optimiser,
scheduler,
save_model_path,
)
if eval_model:
if save_model_path.exists():
model.load_state_dict(torch.load(str(save_model_path)))
print("loading saved model")
with TorchDeviceSession(global_torch_device(cuda_if_available=False),
model):
with torch.no_grad():
with TorchCacheSession():
with TorchEvalSession(model):
valid_masks = []
a = (350, 525)
tr = min(len(valid_loader.dataset) * 4, 2000)
probabilities = numpy.zeros((tr, *a),
dtype=numpy.float32)
for sample_i, (data, target) in enumerate(
tqdm(valid_loader)):
data = data.to(global_torch_device())
target = target.cpu().detach().numpy()
outpu, *_ = model(data)
outpu = torch.sigmoid(outpu).cpu().detach().numpy()
for p in range(data.shape[0]):
output, mask = outpu[p], target[p]
"""
for m in mask:
valid_masks.append(cv2_resize(m, a))
for probability in output:
probabilities[sample_i, :, :] = cv2_resize(probability, a)
sample_i += 1
"""
if sample_i >= tr - 1:
break
if sample_i >= tr - 1:
break
f, ax = pyplot.subplots(3, 3, figsize=(24, 12))
for i in range(3):
ax[0, i].imshow(valid_masks[i], vmin=0, vmax=1)
ax[0, i].set_title("Original", fontsize=14)
ax[1, i].imshow(valid_masks[i], vmin=0, vmax=1)
ax[1, i].set_title("Target", fontsize=14)
ax[2, i].imshow(probabilities[i], vmin=0, vmax=1)
ax[2, i].set_title("Prediction", fontsize=14)
pyplot.show()
def train_d(
model,
train_loader,
valid_loader,
criterion,
optimiser,
scheduler,
save_model_path,
n_epochs=0,
):
"""
Args:
model:
train_loader:
valid_loader:
criterion:
optimiser:
scheduler:
save_model_path:
n_epochs:
Returns:
"""
valid_loss_min = numpy.Inf # track change in validation loss
E = tqdm(range(1, n_epochs + 1))
for epoch in E:
train_loss = 0.0
valid_loss = 0.0
dice_score = 0.0
with TorchTrainSession(model):
train_set = tqdm(train_loader, postfix={"train_loss": 0.0})
for data, target in train_set:
data, target = (
data.to(global_torch_device()),
target.to(global_torch_device()),
)
optimiser.zero_grad()
output, *_ = model(data)
output = torch.sigmoid(output)
loss = criterion(output, target)
loss.backward()
optimiser.step()
train_loss += loss.item() * data.size(0)
train_set.set_postfix(ordered_dict={"train_loss": loss.item()})
with TorchEvalSession(model):
with torch.no_grad():
validation_set = tqdm(valid_loader,
postfix={
"valid_loss": 0.0,
"dice_score": 0.0
})
for data, target in validation_set:
data, target = (
data.to(global_torch_device()),
target.to(global_torch_device()),
)
# forward pass: compute predicted outputs by passing inputs to the model
output, *_ = model(data)
output = torch.sigmoid(output)
# calculate the batch loss
loss = criterion(output, target)
# update average validation loss
valid_loss += loss.item() * data.size(0)
dice_cof = intersection_over_union(
output.cpu().detach().numpy(),
target.cpu().detach().numpy())
dice_score += dice_cof * data.size(0)
validation_set.set_postfix(ordered_dict={
"valid_loss": loss.item(),
"dice_score": dice_cof
})
# calculate average losses
train_loss = train_loss / len(train_loader.dataset)
valid_loss = valid_loss / len(valid_loader.dataset)
dice_score = dice_score / len(valid_loader.dataset)
# print training/validation statistics
E.set_description(f"Epoch: {epoch}"
f" Training Loss: {train_loss:.6f} "
f"Validation Loss: {valid_loss:.6f} "
f"Dice Score: {dice_score:.6f}")
# save model if validation loss has decreased
if valid_loss <= valid_loss_min:
print(
f"Validation loss decreased ({valid_loss_min:.6f} --> {valid_loss:.6f}). Saving model ..."
)
torch.save(model.state_dict(), save_model_path)
valid_loss_min = valid_loss
scheduler.step()
model, scheduler = reschedule(model, epoch, scheduler)
return model
def main():
dataset_root = Path.home() / "Data"
base_path = ensure_existence(PROJECT_APP_PATH.user_data / 'maskrcnn')
log_path = ensure_existence(PROJECT_APP_PATH.user_log / 'maskrcnn')
export_root = ensure_existence(base_path / 'models')
model_name = f'maskrcnn_pennfudanped'
batch_size = 4
num_epochs = 10
optimiser_spec = GDKC(torch.optim.Adam, lr=3e-4)
scheduler_spec = GDKC(
torch.optim.lr_scheduler.
StepLR, # a learning rate scheduler which decreases the learning rate by
step_size=3, # 10x every 3 epochs
gamma=0.1,
)
num_workers = os.cpu_count()
torch_seed(3825)
dataset = PennFudanDataset(dataset_root / "PennFudanPed",
Split.Training,
return_variant=ReturnVariant.all)
dataset_validation = PennFudanDataset(
dataset_root / "PennFudanPed",
Split.Validation,
return_variant=ReturnVariant.all,
)
split = SplitIndexer(len(dataset), validation=0.3, testing=0)
split_indices = torch.randperm(split.total_num).tolist()
data_loader = DataLoader(
Subset(dataset, split_indices[:-split.validation_num]),
batch_size=batch_size,
shuffle=True,
num_workers=num_workers,
collate_fn=collate_batch_fn,
)
data_loader_val = DataLoader(
Subset(dataset_validation, split_indices[-split.validation_num:]),
batch_size=1,
shuffle=False,
num_workers=num_workers,
collate_fn=collate_batch_fn,
)
model = get_pretrained_instance_segmentation_maskrcnn(
dataset.response_channels)
optimiser = optimiser_spec(trainable_parameters(model))
lr_scheduler = scheduler_spec(optimiser)
if True:
model = load_model(model_name=model_name,
model_directory=export_root)
if True:
with TorchTrainSession(model):
with TensorBoardPytorchWriter(log_path / model_name) as writer:
for epoch_i in tqdm(range(num_epochs), desc="Epoch #"):
maskrcnn_train_single_epoch(model=model,
optimiser=optimiser,
data_loader=data_loader,
writer=writer)
lr_scheduler.step() # update the learning rate
maskrcnn_evaluate(
model, data_loader_val, writer=writer
) # evaluate on the validation dataset
save_model(model,
model_name=model_name,
save_directory=export_root)
if True:
with TorchEvalSession(model): # put the model in evaluation mode
img, _ = dataset_validation[
0] # pick one image from the test set
with torch.no_grad():
prediction = model([img.to(global_torch_device())])
from matplotlib import pyplot
pyplot.imshow(
Image.fromarray(
img.mul(255).permute(1, 2, 0).byte().numpy()))
pyplot.show()
import cv2
pyplot.imshow(
Image.fromarray(prediction[0]["masks"][0, 0].mul(
255).byte().cpu().numpy()))
pyplot.show()
(boxes, labels, scores) = (
prediction[0]["boxes"].to('cpu').numpy(),
prediction[0]["labels"].to('cpu').numpy(),
torch.sigmoid(prediction[0]["scores"]).to('cpu').numpy(),
)
from draugr.opencv_utilities import draw_bounding_boxes
from draugr.torch_utilities.images.conversion import quick_to_pil_image
indices = scores > 0.1
cv2.namedWindow(model_name, cv2.WINDOW_NORMAL)
cv2.imshow(
model_name,
draw_bounding_boxes(
quick_to_pil_image(img),
boxes[indices],
labels=labels[indices],
scores=scores[indices],
#categories=categories,
))
cv2.waitKey()
def export_detection_model(
model_export_path: Path = ensure_existence(
PROJECT_APP_PATH.user_data / "penn_fudan_segmentation"
)
/ "seg_skip_fis",
SEED: int = 87539842,
) -> None:
"""
:param model_export_path:
:type model_export_path:
:return:
:rtype:"""
model = OutputActivationModule(
SkipHourglassFission(input_channels=3, output_heads=(1,), encoding_depth=1)
)
with TorchDeviceSession(device=global_torch_device("cpu"), model=model):
with TorchEvalSession(model):
seed_stack(SEED)
# standard PyTorch mean-std input image normalization
transform = transforms.Compose(
[
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
]
)
frame_g = frame_generator(cv2.VideoCapture(0))
for image in tqdm(frame_g):
example_input = (
transform(image).unsqueeze(0).to(global_torch_device()),
)
try:
traced_script_module = torch.jit.trace(
model,
example_input,
# strict=strict_jit,
check_inputs=(
transform(next(frame_g))
.unsqueeze(0)
.to(global_torch_device()),
transform(next(frame_g))
.unsqueeze(0)
.to(global_torch_device()),
),
)
exp_path = model_export_path.with_suffix(".traced")
traced_script_module.save(str(exp_path))
print(
f"Traced Ops used {torch.jit.export_opnames(traced_script_module)}"
)
sprint(
f"Successfully exported JIT Traced model at {exp_path}",
color="green",
)
except Exception as e_i:
sprint(f"Torch JIT Trace export does not work!, {e_i}", color="red")
break
def run_webcam_demo(
cfg: NOD,
categories: Sequence[str],
model_checkpoint: Path,
score_threshold: float = 0.5,
window_name: str = "SSD",
):
"""
:param categories:
:type categories:
:param cfg:
:type cfg:
:param model_checkpoint:
:type model_checkpoint:
:param score_threshold:
:type score_threshold:
:param window_name:
:type window_name:
:return:
:rtype:"""
cpu_device = torch.device("cpu")
transforms = SSDTransform(cfg.input.image_size,
cfg.input.pixel_mean,
split=SplitEnum.testing)
model = SingleShotDetection(cfg)
checkpointer = CheckPointer(model,
save_dir=ensure_existence(
PROJECT_APP_PATH.user_data / "results"))
checkpointer.load(model_checkpoint, use_latest=model_checkpoint is None)
print(
f"Loaded weights from {model_checkpoint if model_checkpoint else checkpointer.get_checkpoint_file()}"
)
model.post_init()
model.to(global_torch_device())
with TorchEvalSession(model):
for infos in tqdm(DictUnityEnvironment(connect_to_running=True)):
info = next(iter(infos.values()))
new_images = extract_all_cameras(info)
image = next(iter(new_images.values()))[..., :3][..., ::-1]
image = gamma_correct_float_to_byte(image)
result = model(
transforms(image)[0].unsqueeze(0).to(global_torch_device()))[0]
height, width, *_ = image.shape
result["boxes"][:, 0::2] *= width / result["img_width"]
result["boxes"][:, 1::2] *= height / result["img_height"]
(boxes, labels, scores) = (
result["boxes"].to(cpu_device).numpy(),
result["labels"].to(cpu_device).numpy(),
result["scores"].to(cpu_device).numpy(),
)
indices = scores > score_threshold
if show_image(
draw_bounding_boxes(
image,
boxes[indices],
labels=labels[indices],
scores=scores[indices],
categories=categories,
).astype(numpy.uint8),
window_name,
wait=1,
):
break # esc to quit
def run_traced_webcam_demo(
input_cfg: NOD,
categories: List,
score_threshold: float = 0.7,
window_name: str = "SSD",
onnx_exported: bool = False,
):
"""
:param onnx_exported:
:type onnx_exported:
:param input_cfg:
:type input_cfg:
:param categories:
:type categories:
:param score_threshold:
:type score_threshold:
:param window_name:
:type window_name:
:return:
:rtype:
"""
pass
import torch
cpu_device = torch.device("cpu")
transforms = SSDTransform(input_cfg.image_size,
input_cfg.pixel_mean,
split=Split.Testing)
model = None
if onnx_exported:
import onnx
onnx_model = onnx.load("torch_model.onnx")
onnx.checker.check_model(onnx_model)
import onnxruntime
ort_session = onnxruntime.InferenceSession("torch_model.onnx")
def to_numpy(tensor):
return (tensor.detach().cpu().numpy()
if tensor.requires_grad else tensor.cpu().numpy())
x = None
# compute onnxruntime output prediction
ort_inputs = {ort_session.get_inputs()[0].name: to_numpy(x)}
ort_outs = ort_session.run(None, ort_inputs)
else:
import torch
import io
torch.jit.load("torch_model.traced")
with open("torch_model.traced",
"rb") as f: # Load ScriptModule from io.BytesIO object
buffer = io.BytesIO(f.read())
model = torch.jit.load(
buffer) # Load all tensors to the original device
"""
buffer.seek(0)
torch.jit.load(buffer, map_location=torch.device('cpu')) # Load all tensors onto CPU, using a device
buffer.seek(0)
model = torch.jit.load(buffer, map_location='cpu') # Load all tensors onto CPU, using a string
# Load with extra files.
extra_files = torch._C.ExtraFilesMap()
extra_files['foo.txt'] = 'bar'
torch.jit.load('torch_model.traced', _extra_files=extra_files)
print(extra_files['foo.txt'])
#exit(0)
"""
with TorchDeviceSession(
device=global_torch_device(cuda_if_available=False), model=model):
with TorchEvalSession(model):
for image in tqdm(frame_generator(cv2.VideoCapture(0))):
result = SSDOut(*model(
transforms(image)[0].unsqueeze(0).to(
global_torch_device())))
height, width, *_ = image.shape
result.boxes[:, 0::2] *= width / result.img_width.cpu().item()
result.boxes[:,
1::2] *= height / result.img_height.cpu().item()
(boxes, labels, scores) = (
result.boxes.to(cpu_device).numpy(),
result.labels.to(cpu_device).numpy(),
result.scores.to(cpu_device).numpy(),
)
indices = scores > score_threshold
cv2.namedWindow(window_name, cv2.WINDOW_NORMAL)
cv2.imshow(
window_name,
draw_bounding_boxes(
image,
boxes[indices],
labels=labels[indices],
scores=scores[indices],
categories=categories,
score_font=ImageFont.truetype(
PACKAGE_DATA_PATH / "Lato-Regular.ttf",
24,
),
).astype(numpy.uint8),
)
if cv2.waitKey(1) == 27:
break # esc to quit
def train_siamese(
model: Module,
optimiser: Optimizer,
criterion: callable,
*,
writer: Writer = MockWriter(),
train_number_epochs: int,
data_dir: Path,
train_batch_size: int,
model_name: str,
save_path: Path,
save_best: bool = False,
img_size: Tuple[int, int],
validation_interval: int = 1,
):
"""
:param img_size:
:type img_size:
:param validation_interval:
:type validation_interval:
:param data_dir:
:type data_dir:
:param optimiser:
:type optimiser:
:param criterion:
:type criterion:
:param writer:
:type writer:
:param model_name:
:type model_name:
:param save_path:
:type save_path:
:param save_best:
:type save_best:
:param model:
:type model:
:param train_number_epochs:
:type train_number_epochs:
:param train_batch_size:
:type train_batch_size:
:return:
:rtype:
"""
train_dataloader = DataLoader(
PairDataset(
data_path=data_dir,
transform=transforms.Compose([
transforms.Grayscale(),
transforms.Resize(img_size),
transforms.ToTensor(),
]),
split=Split.Training,
),
shuffle=True,
num_workers=4,
batch_size=train_batch_size,
)
valid_dataloader = DataLoader(
PairDataset(
data_path=data_dir,
transform=transforms.Compose([
transforms.Grayscale(),
transforms.Resize(img_size),
transforms.ToTensor(),
]),
split=Split.Validation,
),
shuffle=True,
num_workers=4,
batch_size=train_batch_size,
)
best = math.inf
E = tqdm(range(0, train_number_epochs))
batch_counter = count()
for epoch in E:
for tss in train_dataloader:
batch_i = next(batch_counter)
with TorchTrainSession(model):
o = [t.to(global_torch_device()) for t in tss]
optimiser.zero_grad()
loss_contrastive = criterion(model(*o[:2]),
o[2].to(dtype=torch.float))
loss_contrastive.backward()
optimiser.step()
train_loss = loss_contrastive.cpu().item()
writer.scalar("train_loss", train_loss, batch_i)
if batch_counter.__next__() % validation_interval == 0:
with TorchEvalSession(model):
for tsv in valid_dataloader:
ov = [t.to(global_torch_device()) for t in tsv]
v_o, fact = model(*ov[:2]), ov[2].to(dtype=torch.float)
valid_loss = criterion(v_o, fact).cpu().item()
valid_accuracy = (accuracy(distances=v_o,
is_diff=fact).cpu().item())
writer.scalar("valid_loss", valid_loss, batch_i)
if valid_loss < best:
best = valid_loss
print(f"new best {best}")
writer.blip("new_best", batch_i)
if save_best:
save_model_parameters(
model,
optimiser=optimiser,
model_name=model_name,
save_directory=save_path,
)
E.set_description(
f"Epoch number {epoch}, Current train loss {train_loss}, valid loss {valid_loss}, valid_accuracy {valid_accuracy}"
)
return model
