How to train non quantized layers of quantized model on GPU

quantization
1

Hi,

I am trying to use resnet50 quantized model for transfer learning and training the last fully connected layer. From the documentation what I figured was quantized model can only be trained on CPUs but the non-quantized layer can be trained on GPU but how to transfer this layer to GPU? Cannot find much related to this anywhere.

Model looks like:

def create_combined_model(model_fe):
  # Step 1. Isolate the feature extractor.
  model_fe_features = nn.Sequential(
    model_fe.quant,  # Quantize the input
    model_fe.conv1,
    model_fe.bn1,
    model_fe.relu,
    model_fe.maxpool,
    model_fe.layer1,
    model_fe.layer2,
    model_fe.layer3,
    model_fe.layer4,
    model_fe.avgpool,
    model_fe.dequant,  # Dequantize the output
  )

  # Step 2. Create a new "head"
  new_head = nn.Sequential(
    nn.Dropout(p=0.5),
    nn.Linear(num_ftrs, num_classes),
  )

  new_model = nn.Sequential(
    model_fe_features,
    nn.Flatten(1),
    new_head,
  )
  return new_model
2

you can move a layer to GPU by: model.layer.cuda() I think? it probably won’t work if you have mixed cpu/cuda layers I feel