[BERT] Switch to use CheckpointReader to load tensors from checkpoint. (#554)

* [BERT] Switch to use CheckpointReader to load tensors from checkpoint (it was using TensorFlow API directly). Also add a simple command line parameter for the example app, BERT-CoLA, which takes "BERT", "ALBERT", "RoBERTa" to switch between models.

* Add `BERTCheckpointReader.swift` into `CMakeLists.txt`. Revert a change to load checkpoints of a different format from TensorFlow Hub. Fix indentations.

* Make BERT-CoLA display correct model name.

* Pass URL to a checkpoint into BERT to create CheckpointReader and then load tensors from it.

* Clean BERT-CoLA/main.swift

Co-authored-by: Brad Larson <[email protected]>
Co-authored-by: Xihui Wu <[email protected]>

Author: 3 people

Examples/BERT-CoLA/main.swift

@@ -21,13 +21,20 @@ import x10_optimizers_optimizer
 
 let device = Device.defaultXLA
 
-let bertPretrained = BERT.PreTrainedModel.bertBase(cased: false, multilingual: false)
-let workspaceURL = URL(
-    fileURLWithPath: "bert_models", isDirectory: true,
-    relativeTo: URL(
-        fileURLWithPath: NSTemporaryDirectory(),
-        isDirectory: true))
-let bert = try BERT.PreTrainedModel.load(bertPretrained)(from: workspaceURL)
+var bertPretrained: BERT.PreTrainedModel
+if CommandLine.arguments.count >= 2 {
+    if CommandLine.arguments[1].lowercased() == "albert" {
+        bertPretrained = BERT.PreTrainedModel.albertBase
+    } else if CommandLine.arguments[1].lowercased() == "roberta" {
+        bertPretrained = BERT.PreTrainedModel.robertaBase
+    } else {
+        bertPretrained = BERT.PreTrainedModel.bertBase(cased: false, multilingual: false)
+    }
+} else {
+    bertPretrained = BERT.PreTrainedModel.bertBase(cased: false, multilingual: false)
+}
+
+let bert = try bertPretrained.load()
 var bertClassifier = BERTClassifier(bert: bert, classCount: 1)
 bertClassifier.move(to: device)
 
@@ -48,6 +55,9 @@ let epochCount = 3
 let stepsPerEpoch = 1068 // function of training set size and batching configuration
 let peakLearningRate: Float = 2e-5
 
+let workspaceURL = URL(fileURLWithPath: "bert_models", isDirectory: true,
+    relativeTo: URL(fileURLWithPath: NSTemporaryDirectory(),isDirectory: true))
+
 var cola = try CoLA(
   taskDirectoryURL: workspaceURL,
   maxSequenceLength: maxSequenceLength,
@@ -85,7 +95,7 @@ var scheduledLearningRate = LinearlyDecayedParameter(
   startStep: 10
 )
 
-print("Training BERT for the CoLA task!")
+print("Training \(bertPretrained.name) for the CoLA task!")
 for (epoch, epochBatches) in cola.trainingEpochs.prefix(3).enumerated() {
     print("[Epoch \(epoch + 1)]")
     Context.local.learningPhase = .training

Models/Text/BERT.swift

@@ -655,39 +655,37 @@ extension BERT {
             }
         }
 
-        /// Loads this pre-trained BERT model from the specified directory.
+        /// Loads this pre-trained BERT model from the specified URL.
         ///
         /// - Note: This function will download the pre-trained model files to the specified
         //    directory, if they are not already there.
         ///
         /// - Parameters:
-        ///   - directory: Directory to load the pretrained model from.
-        public func load(from directory: URL) throws -> BERT {
+        ///   - url: URL to load the pretrained model from.
+        public func load(from url: URL? = nil) throws -> BERT {
             print("Loading BERT pre-trained model '\(name)'.")
-            let directory = directory.appendingPathComponent(variant.description, isDirectory: true)
-            try maybeDownload(to: directory)
+            
+            let reader = try CheckpointReader(checkpointLocation: url ?? self.url, modelName: name)
+            // TODO(michellecasbon): expose this.
+            reader.isCRCVerificationEnabled = false
+
+            let storage = reader.localCheckpointLocation.deletingLastPathComponent()
 
             // Load the appropriate vocabulary file.
             let vocabulary: Vocabulary = {
                 switch self {
                 case .bertBase, .bertLarge:
-                    let vocabularyURL = directory
-                        .appendingPathComponent(subDirectory)
-                        .appendingPathComponent("vocab.txt")
+                    let vocabularyURL = storage.appendingPathComponent("vocab.txt")
                     return try! Vocabulary(fromFile: vocabularyURL)
                 case .robertaBase, .robertaLarge:
-                    let vocabularyURL = directory
-                        .appendingPathComponent(subDirectory)
-                        .appendingPathComponent("vocab.json")
-                    let dictionaryURL = directory
-                        .appendingPathComponent(subDirectory)
-                        .appendingPathComponent("dict.txt")
+                    let vocabularyURL = storage.appendingPathComponent("vocab.json")
+                    let dictionaryURL = storage.appendingPathComponent("dict.txt")
                     return try! Vocabulary(
                         fromRoBERTaJSONFile: vocabularyURL,
                         dictionaryFile: dictionaryURL)
                 case .albertBase, .albertLarge, .albertXLarge, .albertXXLarge:
-                    let vocabularyURL = directory
-                        .appendingPathComponent(subDirectory)
+                    let vocabularyURL = storage
+                        .deletingLastPathComponent()
                         .appendingPathComponent("assets")
                         .appendingPathComponent("30k-clean.model")
                     return try! Vocabulary(fromSentencePieceModel: vocabularyURL)
@@ -704,8 +702,7 @@ extension BERT {
                         unknownToken: "[UNK]",
                         maxTokenLength: nil)
                 case .robertaBase, .robertaLarge:
-                    let mergePairsFileURL = directory
-                        .appendingPathComponent(subDirectory)
+                    let mergePairsFileURL = storage
                         .appendingPathComponent("merges.txt")
                     let mergePairs = [BytePairEncoder.Pair: Int](
                         uniqueKeysWithValues:
@@ -749,150 +746,40 @@ extension BERT {
                 initializerStandardDeviation: 0.02,
                 useOneHotEmbeddings: false)
 
-            // Load the pre-trained model checkpoint.
-            switch self {
-            case .bertBase, .bertLarge:
-                model.load(fromTensorFlowCheckpoint: directory
-                    .appendingPathComponent(subDirectory)
-                    .appendingPathComponent("bert_model.ckpt"))
-            case .robertaBase, .robertaLarge:
-                model.load(fromTensorFlowCheckpoint: directory
-                    .appendingPathComponent(subDirectory)
-                    .appendingPathComponent("roberta_\(subDirectory).ckpt"))
-            case .albertBase, .albertLarge, .albertXLarge, .albertXXLarge:
-                model.load(fromTensorFlowCheckpoint: directory
-                    .appendingPathComponent(subDirectory)
-                    .appendingPathComponent("variables")
-                    .appendingPathComponent("variables"))
-            }
+            model.loadTensors(reader)
             return model
         }
-
-        /// Downloads this pre-trained model to the specified directory, if it's not already there.
-        public func maybeDownload(to directory: URL) throws {
-            switch self {
-            case .bertBase, .bertLarge, .robertaBase, .robertaLarge:
-                // Download and extract the pretrained model, if necessary.
-                DatasetUtilities.downloadResource(filename: "\(subDirectory)", fileExtension: "zip",
-                                                  remoteRoot: url.deletingLastPathComponent(),
-                                                  localStorageDirectory: directory)
-            case .albertBase, .albertLarge, .albertXLarge, .albertXXLarge:
-                // Download the model, if necessary.
-                let compressedFileURL = directory.appendingPathComponent("\(subDirectory).tar.gz")
-                try download(from: url, to: compressedFileURL)
-
-                // Extract the data, if necessary.
-                let extractedDirectoryURL = directory.appendingPathComponent(subDirectory)
-                if !FileManager.default.fileExists(atPath: extractedDirectoryURL.path) {
-                    try extract(tarGZippedFileAt: compressedFileURL, to: extractedDirectoryURL)
-                }
-            }
-        }
   }
 
-    /// Loads a BERT model from the provided TensorFlow checkpoint file into this BERT model.
+    /// Loads a BERT model from the provided CheckpointReader into this BERT model.
     ///
     /// - Parameters:
-    ///   - fileURL: Path to the checkpoint file. Note that TensorFlow checkpoints typically
-    ///     consist of multiple files (e.g., `bert_model.ckpt.index`, `bert_model.ckpt.meta`, and
-    ///     `bert_model.ckpt.data-00000-of-00001`). In this case, the file URL should be specified
-    ///     as their common prefix (e.g., `bert_model.ckpt`).
-    public mutating func load(fromTensorFlowCheckpoint fileURL: URL) {
-        let checkpointReader = TensorFlowCheckpointReader(checkpointPath: fileURL.path)
-        tokenEmbedding.embeddings =
-            Tensor(checkpointReader.loadTensor(named: "bert/embeddings/word_embeddings"))
-        positionEmbedding.embeddings =
-            Tensor(checkpointReader.loadTensor(named: "bert/embeddings/position_embeddings"))
-        embeddingLayerNorm.offset =
-            Tensor(checkpointReader.loadTensor(named: "bert/embeddings/LayerNorm/beta"))
-        embeddingLayerNorm.scale =
-            Tensor(checkpointReader.loadTensor(named: "bert/embeddings/LayerNorm/gamma"))
+    ///   - reader: CheckpointReader object to load tensors from.
+    public mutating func loadTensors(_ reader: CheckpointReader) {
+        tokenEmbedding.embeddings = reader.readTensor(name: "bert/embeddings/word_embeddings")
+        positionEmbedding.embeddings = reader.readTensor(name: "bert/embeddings/position_embeddings")
+        embeddingLayerNorm.offset = reader.readTensor(name: "bert/embeddings/LayerNorm/beta")
+        embeddingLayerNorm.scale = reader.readTensor(name: "bert/embeddings/LayerNorm/gamma")
         switch variant {
         case .bert, .albert:
             tokenTypeEmbedding.embeddings =
-                Tensor(checkpointReader.loadTensor(named: "bert/embeddings/token_type_embeddings"))
+                reader.readTensor(name: "bert/embeddings/token_type_embeddings")
         case .roberta: ()
         }
         switch variant {
         case .bert, .roberta:
             for layerIndex in encoderLayers.indices {
-                let prefix = "bert/encoder/layer_\(layerIndex)"
-                encoderLayers[layerIndex].multiHeadAttention.queryWeight =
-                    Tensor(checkpointReader.loadTensor(named: "\(prefix)/attention/self/query/kernel"))
-                encoderLayers[layerIndex].multiHeadAttention.queryBias =
-                    Tensor(checkpointReader.loadTensor(named: "\(prefix)/attention/self/query/bias"))
-                encoderLayers[layerIndex].multiHeadAttention.keyWeight =
-                    Tensor(checkpointReader.loadTensor(named: "\(prefix)/attention/self/key/kernel"))
-                encoderLayers[layerIndex].multiHeadAttention.keyBias =
-                    Tensor(checkpointReader.loadTensor(named: "\(prefix)/attention/self/key/bias"))
-                encoderLayers[layerIndex].multiHeadAttention.valueWeight =
-                    Tensor(checkpointReader.loadTensor(named: "\(prefix)/attention/self/value/kernel"))
-                encoderLayers[layerIndex].multiHeadAttention.valueBias =
-                    Tensor(checkpointReader.loadTensor(named: "\(prefix)/attention/self/value/bias"))
-                encoderLayers[layerIndex].attentionWeight =
-                    Tensor(checkpointReader.loadTensor(named: "\(prefix)/attention/output/dense/kernel"))
-                encoderLayers[layerIndex].attentionBias =
-                    Tensor(checkpointReader.loadTensor(named: "\(prefix)/attention/output/dense/bias"))
-                encoderLayers[layerIndex].attentionLayerNorm.offset =
-                    Tensor(checkpointReader.loadTensor(named: "\(prefix)/attention/output/LayerNorm/beta"))
-                encoderLayers[layerIndex].attentionLayerNorm.scale =
-                    Tensor(checkpointReader.loadTensor(named: "\(prefix)/attention/output/LayerNorm/gamma"))
-                encoderLayers[layerIndex].intermediateWeight =
-                    Tensor(checkpointReader.loadTensor(named: "\(prefix)/intermediate/dense/kernel"))
-                encoderLayers[layerIndex].intermediateBias =
-                    Tensor(checkpointReader.loadTensor(named: "\(prefix)/intermediate/dense/bias"))
-                encoderLayers[layerIndex].outputWeight =
-                    Tensor(checkpointReader.loadTensor(named: "\(prefix)/output/dense/kernel"))
-                encoderLayers[layerIndex].outputBias =
-                    Tensor(checkpointReader.loadTensor(named: "\(prefix)/output/dense/bias"))
-                encoderLayers[layerIndex].outputLayerNorm.offset =
-                    Tensor(checkpointReader.loadTensor(named: "\(prefix)/output/LayerNorm/beta"))
-                encoderLayers[layerIndex].outputLayerNorm.scale =
-                    Tensor(checkpointReader.loadTensor(named: "\(prefix)/output/LayerNorm/gamma"))
+                encoderLayers[layerIndex].load(bert: reader,
+                    prefix: "bert/encoder/layer_\(layerIndex)")
             }
         case .albert:
             embeddingProjection[0].weight =
-                Tensor(checkpointReader.loadTensor(
-                    named: "bert/encoder/embedding_hidden_mapping_in/kernel"))
+                reader.readTensor(name: "bert/encoder/embedding_hidden_mapping_in/kernel")
             embeddingProjection[0].bias =
-                Tensor(checkpointReader.loadTensor(
-                    named: "bert/encoder/embedding_hidden_mapping_in/bias"))
+                reader.readTensor(name: "bert/encoder/embedding_hidden_mapping_in/bias")
             for layerIndex in encoderLayers.indices {
                 let prefix = "bert/encoder/transformer/group_\(layerIndex)/inner_group_0"
-                encoderLayers[layerIndex].multiHeadAttention.queryWeight =
-                    Tensor(checkpointReader.loadTensor(named: "\(prefix)/attention_1/self/query/kernel"))
-                encoderLayers[layerIndex].multiHeadAttention.queryBias =
-                    Tensor(checkpointReader.loadTensor(named: "\(prefix)/attention_1/self/query/bias"))
-                encoderLayers[layerIndex].multiHeadAttention.keyWeight =
-                    Tensor(checkpointReader.loadTensor(named: "\(prefix)/attention_1/self/key/kernel"))
-                encoderLayers[layerIndex].multiHeadAttention.keyBias =
-                    Tensor(checkpointReader.loadTensor(named: "\(prefix)/attention_1/self/key/bias"))
-                encoderLayers[layerIndex].multiHeadAttention.valueWeight =
-                    Tensor(checkpointReader.loadTensor(named: "\(prefix)/attention_1/self/value/kernel"))
-                encoderLayers[layerIndex].multiHeadAttention.valueBias =
-                    Tensor(checkpointReader.loadTensor(named: "\(prefix)/attention_1/self/value/bias"))
-                encoderLayers[layerIndex].attentionWeight =
-                    Tensor(checkpointReader.loadTensor(named: "\(prefix)/attention_1/output/dense/kernel"))
-                encoderLayers[layerIndex].attentionBias =
-                    Tensor(checkpointReader.loadTensor(named: "\(prefix)/attention_1/output/dense/bias"))
-                encoderLayers[layerIndex].attentionLayerNorm.offset =
-                    Tensor(checkpointReader.loadTensor(named: "\(prefix)/LayerNorm/beta"))
-                encoderLayers[layerIndex].attentionLayerNorm.scale =
-                    Tensor(checkpointReader.loadTensor(named: "\(prefix)/LayerNorm/gamma"))
-                encoderLayers[layerIndex].intermediateWeight =
-                    Tensor(checkpointReader.loadTensor(named: "\(prefix)/ffn_1/intermediate/dense/kernel"))
-                encoderLayers[layerIndex].intermediateBias =
-                    Tensor(checkpointReader.loadTensor(named: "\(prefix)/ffn_1/intermediate/dense/bias"))
-                encoderLayers[layerIndex].outputWeight =
-                    Tensor(checkpointReader.loadTensor(
-                        named: "\(prefix)/ffn_1/intermediate/output/dense/kernel"))
-                encoderLayers[layerIndex].outputBias =
-                    Tensor(checkpointReader.loadTensor(
-                        named: "\(prefix)/ffn_1/intermediate/output/dense/bias"))
-                encoderLayers[layerIndex].outputLayerNorm.offset =
-                    Tensor(checkpointReader.loadTensor(named: "\(prefix)/LayerNorm_1/beta"))
-                encoderLayers[layerIndex].outputLayerNorm.scale =
-                    Tensor(checkpointReader.loadTensor(named: "\(prefix)/LayerNorm_1/gamma"))
+                encoderLayers[layerIndex].load(albert: reader, prefix: prefix)
             }
         }
     }