bot commited on
Commit ·
44038e9
1
Parent(s): 3f121e9
Auto-detect resume from checkpoint, save_freq 5000
Browse files- train_cloud.sh +12 -1
train_cloud.sh
CHANGED
|
@@ -29,10 +29,21 @@ if [ "$NUM_GPUS" -gt 1 ]; then
|
|
| 29 |
ACCEL_FLAGS="--multi_gpu --num_processes $NUM_GPUS"
|
| 30 |
fi
|
| 31 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 32 |
python3.12 -m accelerate.commands.launch $ACCEL_FLAGS \
|
| 33 |
-m lerobot.scripts.lerobot_train \
|
|
|
|
| 34 |
--dataset.repo_id="so100:$DATASET_DIR:/workspace/pi05-so100-diverse/filtered_index.json:/workspace/pi05-so100-diverse/norm_stats.json" \
|
| 35 |
-
--policy.path=lerobot/pi05_base \
|
| 36 |
--policy.train_expert_only=true \
|
| 37 |
--policy.dtype=bfloat16 \
|
| 38 |
--policy.gradient_checkpointing=false \
|
|
|
|
| 29 |
ACCEL_FLAGS="--multi_gpu --num_processes $NUM_GPUS"
|
| 30 |
fi
|
| 31 |
|
| 32 |
+
# Check if we can resume from a checkpoint
|
| 33 |
+
RESUME_ARGS=""
|
| 34 |
+
LAST_CKPT="/ephemeral/production_run/checkpoints/last/pretrained_model/train_config.json"
|
| 35 |
+
if [ -f "$LAST_CKPT" ]; then
|
| 36 |
+
echo "Resuming from checkpoint: $LAST_CKPT" | tee -a "$LOG_FILE"
|
| 37 |
+
RESUME_ARGS="--resume=true --config_path=$LAST_CKPT"
|
| 38 |
+
else
|
| 39 |
+
echo "Starting fresh training" | tee -a "$LOG_FILE"
|
| 40 |
+
RESUME_ARGS="--policy.path=lerobot/pi05_base"
|
| 41 |
+
fi
|
| 42 |
+
|
| 43 |
python3.12 -m accelerate.commands.launch $ACCEL_FLAGS \
|
| 44 |
-m lerobot.scripts.lerobot_train \
|
| 45 |
+
$RESUME_ARGS \
|
| 46 |
--dataset.repo_id="so100:$DATASET_DIR:/workspace/pi05-so100-diverse/filtered_index.json:/workspace/pi05-so100-diverse/norm_stats.json" \
|
|
|
|
| 47 |
--policy.train_expert_only=true \
|
| 48 |
--policy.dtype=bfloat16 \
|
| 49 |
--policy.gradient_checkpointing=false \
|