Memory error - Optimization to increase batch size

When training with samples of size 256x256 pixels, a batch size over 32 causes a memory error from cuda. We have to find a way to optimize the process, in order to increase the batch size.

RunTimeError: cuda runtime error (2) : out of memory at /opt/conda/ .../THCStorage.cu:58

NOTE : May be specific to our (GC HPC) computing environment

[epeterson12]

Validation of Checkpointed results

Results when using checkpointing are slightly different from those of the original unetsmall model because CudNN has non-deterministic kernels. I ran tests using suggestions from pytorch discussions https://discuss.pytorch.org/t/non-reproducible-result-with-gpu/1831 and https://discuss.pytorch.org/t/deterministic-non-deterministic-results-with-pytorch/9087.

Using the same sample files, I ran train_model.py twice using the unetsmall model (batch_size = 32) and I ran it once using the checkpointed_unet model (batch_size = 50). Then, I classified some images with the resulting models.

The settings to try to get reproducible results were set as follows at the beginning of the code:

torch.backends.cudnn.deterministic = True
torch.manual_seed(999)
torch.cuda.manual_seed(999)
torch.cuda.manual_seed_all(999)
random.seed(0)

Also, in the DataLoaders, the parameters used for the instantiation had num_workers = 0 and shuffle = False Running the original unetsmall configuration without checkpoints 2 times yielded two slightly different results. Here are some examples of the results obtained when running image_classification.py on one of the training images with each trained model. Sections if the images that weren’t classified were left white.

Ground Truth	Current Code	Current Code 2	Checkpointed

Please note that the configurations and the number of samples weren’t set to yield optimal results. Verifying reproducibility was the goal of these tests. The number of training samples was set to the number of samples produced during the samples creation.

global:
  samples_size: 256
  num_classes: 5
  data_path: /my/data/path
  number_of_bands: 3
  model_name: unetsmall     # One of unet, unetsmall, checkpointed_unet or ternausnet

sample:
  prep_csv_file: /my/prep/csv/file
  samples_dist: 200
  remove_background: True
  mask_input_image: False

training:
  output_path: /my/output/path
  num_trn_samples: 3356
  num_val_samples: 1370
  batch_size: 32
  num_epochs: 100
  learning_rate: 0.0001
  weight_decay: 0
  step_size: 4
  gamma: 0.9
  class_weights: False

models:
  unet:   &unet001
    dropout: False
    probability: 0.2    # Set with dropout
    pretrained: False   # optional
  unetsmall:
    <<: *unet001
  ternausnet:
    pretrained: ./models/TernausNet.pt    # Mandatory
  checkpointed_unet: 
    <<: *unet001

I think that the results of the checkpointed_unet are similar enough to the unetsmall’s results for us to consider that it is a good memory and time optimised version of our unetsmall net architecture. I have added it as a model choice for our program.

Throughout my tests, I observed that the models produced by training are more accurate when the randoms aren’t seeded. The checkpointed_unet, observationally, seems to be more affected by this then the unetsmall.

[epeterson12]

Using checkpointing in the unetsmall net increases the speed of training. Tests were performed using the following parameters:

# Training Samples	# Validation Samples	Sample Size	# Classes	# Epochs
781	495	256	11	200

Learning Rate	Weight Decay	Step Size	Gamma	Class Weights	Dropout
0.0001	0	4	0.9	False	False

Best results

.	Original	Checkpointed
Max batch size	32	50
Time to complete training over 200 epochs	350 min	316 min

Using checkpoints in the net design does seem to affect the results of the training. Tests were done on the original and on the checkpointed nets while setting the random seed to 7 and the models outputted gave similar but slightly different results. In the first test, the original algorithm gave results closer to the ground truth. In the second tests, the checkpointed version of the net yielded better results.