diff --git a/.drone.yml b/.drone.yml index 2ca7c02..2a26d66 100644 --- a/.drone.yml +++ b/.drone.yml @@ -20,8 +20,8 @@ steps: - pip install -q torchsummary matplotlib pytorch-metric-learning torchinfo torch torchvision tqdm Pillow facenet-pytorch - apt update -qq - apt install fonts-noto-cjk-extra -y -qq - - ls ./ + # - ls ./ - mkdir -p data - $mount_command - - ls data/ + # - ls data/ - CI=True python finetune/facenet_transfer_learning.py \ No newline at end of file diff --git a/finetune/facenet_transfer_learning.py b/finetune/facenet_transfer_learning.py index b543b5b..bf2eedf 100644 --- a/finetune/facenet_transfer_learning.py +++ b/finetune/facenet_transfer_learning.py @@ -158,13 +158,11 @@ for epoch in range(epochs): label_text=image_folder['train'].classes) image_pallets.save(join(save_dir, 'pallets', str(epoch) + '_train.jpg')) optimizer.zero_grad() - with autocast(dtype=float16): + with autocast(dtype=float16,enabled=True): images = images.to(device) - labels = labels.to(device) - - outputs = model(images) - - loss = criterion(outputs, labels) + labels = labels.to(device) + outputs = model(images) + loss = criterion(outputs, labels) train_loss += loss.item() scaler.scale(loss).backward() @@ -187,8 +185,9 @@ for epoch in range(epochs): image_pallets = plot_dataset(dataloader=(images, labels), col_len=6, label_text=image_folder['train'].classes) image_pallets.save(join(save_dir, 'pallets', str(epoch) + '_val.jpg')) - images = images.to(device) - labels = labels.to(device) + with autocast(dtype=float16, enabled=True): + images = images.to(device) + labels = labels.to(device) outputs = model_gpu(images) loss = criterion(outputs, labels) val_loss += loss.item()