Skip to content

Train a Segmentation Model

haltosan edited this page Apr 18, 2023 · 4 revisions

Training on Collab

One of the benefits of using Google Colab is that it helps us visualize how our model is doing and it provides free compute/GPU time. This notebook is setup to help us train models from LabelMe or Coco datasets:

https://colab.research.google.com/drive/1pKUldyaq9FiH25LDJXQLeHEVJVhqzWBP?usp=sharing https://colab.research.google.com/drive/17GtCt0uinU5sm0fE1YdsxpdKvahfkYEm?usp=sharing

Click on the link and select "File", "Save a copy in Drive" at the top left corner

Step 1: Load Labeled Images to Google Drive

You can do this by selecting "New", "Upload Folder" in the top left corner of your Google Drive and selecting the labeled images folder.

Step 2: Open and Modify Code

Open the notebook (same as above). Make sure you are editing your copy of it (after clicking 'save a copy in drive') as opposed to the template.

Change the name of code

Click on title in the top left corner and type in the new name you would like.

Set project parameters

Before

# file structure (assuming the path to everyting is at 'My Drive/work/PROJECT_DIR/')
project_dir = 'coloradoLandHWR'
image_relative_folder_path = 'row training data 2'

# model input path; if none, read the model from the default output path 
model_input_path:str = None

# model infomation 
model:str = 'COCO-Detection/faster_rcnn_R_50_FPN_3x.yaml'  # can be COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml or COCO-Detection/faster_rcnn_R_50_FPN_3x.yaml; the faster one doesn't need a mask (some datasets don't have a mask)
num_classes:int = 1  # how many different boxes are there in your training data (if you have first name, last name, birth date => 3 classes)

After

# file structure (assuming the path to everyting is at 'My Drive/work/PROJECT_DIR/')
project_dir = 'mexicoCensus'
image_relative_folder_path = 'labeled images'

# model input path; if none, read the model from the default output path 
model_input_path:str = None

# model infomation 
model:str = 'COCO-Detection/faster_rcnn_R_50_FPN_3x.yaml'  # can be COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml or COCO-Detection/faster_rcnn_R_50_FPN_3x.yaml; the faster one doesn't need a mask (some datasets don't have a mask)
num_classes:int = 12  # how many different boxes are there in your training data (if you have first name, last name, birth date => 3 classes)

Step 3: Run the code

Follow the notebook. There are sections that go over setting up the runtime, setting up the dataset, setting up the model, training, and segmentation. The text cells between code cells should describe what each group of cells do. Comments within each cell can give more specific guidance / explanations. You'll need to stop 2 times while setting up the runtime: once to restart for a dependency and once for mounting your drive. After running the first cell, you'll get a message like this:

  [pydevd_plugins]
You must restart the runtime in order to use newly installed version.
[RESTART RUNTIME]

Click the 'RESTART RUNTIME' button before continuing.

When mounting your drive, you'll need to sign in and allow Google Collab access to your drive. This allows you to read and write files directly from your online storage.

The block that runs the labelme2coco.convert(labelme_folder, save_json_path) statement tends to take a long-ish time for large datasets. If it takes too long (> 30 minutes) stop that cell and rerun.

Some blocks will give you visual output. This allows you to gauge if the images are loading correctly, labeling is working, etc.

Training is the most important part of this notebook. Depending on the size of your images and the number of images or classes, this chunk of code could take a few hours to finish. Do not let your computer fall asleep during this. Google Collab will disconnect your runtime if the notebook is inactive for too long. Make sure you click around and/or edit something useless (like a comment). Otherwise, it might not finish executing.

Once the code is done running, continue to run the code until you finish the next set of visual output code. This shows you how well the segmenting did on the images.

Step 4: Download the weights

Run the prediction chunks of code to get some visualization on how the segmentation might look. (Keep in mind, this might be more accurate than what will be done on the supercomputer because you are running the model on your training dataset.) If you like how the model looks, you can download the weights for the model from the file section on the left side of Google Colab. Click the 3 dots to the left of "model_final.pth" and select download. This might take a while.

Step 5: Move weights to the supercomputer

You will need to use scp from your command line in order to move the weights over to the supercomputer. In the command line, move to the location of your weights, (probably downloads). Then use scp to copy over the weights

It will look something like this:

scp weights_name.pth username@ssh.rc.byu.edu:/fslhome/username/fsl_groups/fslg_groupname/compute/desired_location/of_weights

Example:

scp death_weights.pth bob@ssh.rc.byu.edu:~/fsl_groups/fslg_census/compute/projects/ohio_death_3

You will need to enter your password and a duo authentication number every time you move something over to the supercomputer.

Congratulations! You have successfully trained a segmentation model!

(Optional): Convert predictions into training data

If your model isn't performing as well as you want, you can convert the output to training data (to further improve the model). This is an example of the process. It converts the model output to LabelMe output. Humans should then go in and validate / modify the output (this normally takes a lot faster than creating the labels from scratch). Once the labels are all edited, this is valid training data for the model.

To edit the model output, see this script.

(Optional): Create snippets on Google Collab (not recommended)

Google Collab gives you limited GPU time and will disconnect your runtime if you leave it running for too long without any "activity". On the other hand, you can always run a small sample of images to gauge the output quality in a more interactive environment. For full datasets, use the supercomputer. If you want to gauge model quality, you can use Google Collab.

There is code, but it is overspecialized for the task it was created for. See the examples at the end of the notebook.