Upload adversarial_training_clip_with_object_token.py

train/adversarial_training_clip_with_object_token.py

@@ -265,7 +265,7 @@ def main(args):
 
     # save final model
     torch.save(unwrap_model(model).model.state_dict(), f'{args.output_dir}/checkpoints/final.pt')
-    torch.save(unwrap_model(proj_head).model.state_dict(), f'{args.output_dir}/checkpoints/final_proj_head.pt')
+    torch.save(unwrap_model(proj_head).state_dict(), f'{args.output_dir}/checkpoints/final_proj_head.pt')
 
     torch.save(optimizer.state_dict(), f'{args.output_dir}/checkpoints/final_opt.pt')
 
@@ -505,13 +505,15 @@ def train_one_epoch(
             wandb.log(log_data)
 
         # save 10 models over the course of training
-        if args.save_checkpoints and (step_total % (args.steps // 10) == 0):
+        if args.save_checkpoints and (step_total % (args.steps // 1) == 0):
             # save model and optimizer state_dict
             torch.save(unwrap_model(model).model.state_dict(), f'{args.output_dir}/checkpoints/step_{step_total}.pt')
+            torch.save(unwrap_model(proj_head).state_dict(), f'{args.output_dir}/checkpoints/step_{step_total}_proj_head.pt')
             torch.save(optimizer.state_dict(), f'{args.output_dir}/checkpoints/step_{step_total}_opt.pt')
         # every 200 steps, save a fallback model, which gets overwritten
-        if step_total % 200 == 0:
+        if step_total % 2 == 0:
             torch.save(unwrap_model(model).model.state_dict(), f'{args.output_dir}/checkpoints/fallback_{step_total}.pt')
+            torch.save(unwrap_model(proj_head).state_dict(), f'{args.output_dir}/checkpoints/fallback_{step_total}_proj_head.pt')
             torch.save(optimizer.state_dict(), f'{args.output_dir}/checkpoints/fallback_{step_total}_opt.pt')
             # remove old fallback models
             for file in os.listdir(f'{args.output_dir}/checkpoints'):