Parallelization failed on tf.keras model

[XiYuan68]

Describe the bug
Attack parallelization failed on tf.keras model

To Reproduce
Steps to reproduce the behavior:

from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense, Embedding, GlobalMaxPooling1D
from textattack.attack_recipes import PWWSRen2019
from textattack import Attacker, AttackArgs


from task import TrainSequence, load_data
from audionlp import get_dataset, get_tokenizer, TwentynewsModelWrapper


data = 'twentynews'
batch_size = 32
epochs = 1
train = TrainSequence(data, batch_size)
validation = load_data(data, 'val')[:-1]
input_shape = train.x.shape[1:]
n_choice = train.y.shape[-1]

model = Sequential()
model.add(Embedding(10000, 128, input_length=500, name='embedding'))
model.add(GlobalMaxPooling1D())
model.add(Dense(n_choice, 'softmax'))
model.summary()
model.compile('adam', 'categorical_crossentropy', ['accuracy'])
model.fit(train, epochs=epochs, validation_data=validation)

dataset = get_dataset(data)
tokenizer = get_tokenizer()
wrapper = TwentynewsModelWrapper(model, tokenizer)
attack = PWWSRen2019.build(wrapper)
attack_args = AttackArgs(num_examples=-1, disable_stdout=True,)
                         # parallel=True, num_workers_per_device=2)
attacker = Attacker(attack, dataset, attack_args)
result = attacker.attack_dataset()

when parallel=False everything is fine

but when parallel=True, num_workers_per_device=2, error raised:

textattack: Running 2 worker(s) on 1 GPU(s).
Traceback (most recent call last):

  File "<ipython-input-6-625f4babcd86>", line 4, in <module>
    result = attacker.attack_dataset()

  File "/home/chenxiyuan/projects/python/TextAttack/textattack/attacker.py", line 427, in attack_dataset
    self._attack_parallel()

  File "/home/chenxiyuan/projects/python/TextAttack/textattack/attacker.py", line 280, in _attack_parallel
    worker_pool = torch.multiprocessing.Pool(

  File "/usr/lib/python3.9/multiprocessing/context.py", line 119, in Pool
    return Pool(processes, initializer, initargs, maxtasksperchild,

  File "/usr/lib/python3.9/multiprocessing/pool.py", line 212, in __init__
    self._repopulate_pool()

  File "/usr/lib/python3.9/multiprocessing/pool.py", line 303, in _repopulate_pool
    return self._repopulate_pool_static(self._ctx, self.Process,

  File "/usr/lib/python3.9/multiprocessing/pool.py", line 326, in _repopulate_pool_static
    w.start()

  File "/usr/lib/python3.9/multiprocessing/process.py", line 121, in start
    self._popen = self._Popen(self)

  File "/usr/lib/python3.9/multiprocessing/context.py", line 284, in _Popen
    return Popen(process_obj)

  File "/usr/lib/python3.9/multiprocessing/popen_spawn_posix.py", line 32, in __init__
    super().__init__(process_obj)

  File "/usr/lib/python3.9/multiprocessing/popen_fork.py", line 19, in __init__
    self._launch(process_obj)

  File "/usr/lib/python3.9/multiprocessing/popen_spawn_posix.py", line 47, in _launch
    reduction.dump(process_obj, fp)

  File "/usr/lib/python3.9/multiprocessing/reduction.py", line 60, in dump
    ForkingPickler(file, protocol).dump(obj)

TypeError: cannot pickle 'weakref' object

by the way, when loading a similar pretrained model containing a custom keras layer like:

class Dropout(Layer):
  def __init__(self, level, **kwargs):
    super(Dropout, self).__init__(**kwargs)
    self.level = level

  def build(self, input_shape):
      pass

  def call(self, inputs):
    return K.dropout(inputs, self.level)

parallelizaion raised a slightly different error:

textattack: Running 2 worker(s) on 1 GPU(s).
Traceback (most recent call last):

  ...

  File "/home/chenxiyuan/projects/python/TextAttack/textattack/attacker.py", line 427, in attack_dataset
    self._attack_parallel()

  File "/home/chenxiyuan/projects/python/TextAttack/textattack/attacker.py", line 280, in _attack_parallel
    worker_pool = torch.multiprocessing.Pool(

  File "/usr/lib/python3.9/multiprocessing/context.py", line 119, in Pool
    return Pool(processes, initializer, initargs, maxtasksperchild,

  File "/usr/lib/python3.9/multiprocessing/pool.py", line 212, in __init__
    self._repopulate_pool()

  File "/usr/lib/python3.9/multiprocessing/pool.py", line 303, in _repopulate_pool
    return self._repopulate_pool_static(self._ctx, self.Process,

  File "/usr/lib/python3.9/multiprocessing/pool.py", line 326, in _repopulate_pool_static
    w.start()

  File "/usr/lib/python3.9/multiprocessing/process.py", line 121, in start
    self._popen = self._Popen(self)

  File "/usr/lib/python3.9/multiprocessing/context.py", line 284, in _Popen
    return Popen(process_obj)

  File "/usr/lib/python3.9/multiprocessing/popen_spawn_posix.py", line 32, in __init__
    super().__init__(process_obj)

  File "/usr/lib/python3.9/multiprocessing/popen_fork.py", line 19, in __init__
    self._launch(process_obj)

  File "/usr/lib/python3.9/multiprocessing/popen_spawn_posix.py", line 47, in _launch
    reduction.dump(process_obj, fp)

  File "/usr/lib/python3.9/multiprocessing/reduction.py", line 60, in dump
    ForkingPickler(file, protocol).dump(obj)

PicklingError: Can't pickle <class 'tensorflow.python.keras.saving.saved_model.load.Dropout'>: attribute lookup Dropout on tensorflow.python.keras.saving.saved_model.load failed

Expected behavior
allow parallelization on tf.keras model

Screenshots or Traceback
If applicable, add screenshots to help explain your problem. Also, copy and paste tracebacks produced by the bug.

System Information (please complete the following information):

• OS: Linux
• Library versions torch==1.7.1, transformers==4.8.2, tensorflow-gpu==2.5.0
• Textattack version: 0.3.0

Additional context
with my limited experience of parallelization, multiprocessing is not very ideal for the job, so maybe consider replacing it with tensorflow APIs?

[qiyanjun]

@sanchit97 can you take a look at this?

[qiyanjun]

@XiYuan68 possible to submit a PR?

[XiYuan68]

sorry I don't know how to fix this...

[sanchit97]

This seems like a serious issue in tf.keras models. One of the possible reasons is that tf.keras models are non-picklable. A short hotfix is here: tensorflow/tensorflow#34697
However, I could not get it to run completely but solved the pickling error. I will fix this once I figure out the multiprocessing error. Thanks for raising this issue @XiYuan68

[sanchit97]

So I was able to get this to work. The main problem lies with the tf.keras object. You have to nest it inside the if __name__ == '__main__'. Here's the code:

import tensorflow as tf
import keras
import pickle
import numpy as np
from textattack.models.wrappers import ModelWrapper
from textattack.datasets import HuggingFaceDataset
from textattack.attack_recipes import PWWSRen2019

import numpy as np
# from keras.utils import to_categorical
# from tensorflow.keras.optimizers import Adam
from tensorflow.keras.utils import to_categorical
from tensorflow.keras import models
from tensorflow.keras import layers
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense
from tensorflow.keras.layers import Flatten
from tensorflow.keras.layers import Dropout
from tensorflow.keras.models import Sequential, load_model, save_model, Model

from nltk.tokenize import word_tokenize, RegexpTokenizer
from textattack import AttackArgs
from textattack.datasets import Dataset
from textattack import Attacker
import tempfile
from tensorflow.python.keras.layers import deserialize, serialize
from tensorflow.python.keras.saving import saving_utils
import torch


NUM_WORDS = 1000
# model = Sequential()
# model.add(Embedding(10000, 128, input_length=500, name='embedding'))
# model.add(GlobalMaxPooling1D())
# model.add(Dense(n_choice, 'softmax'))
# model.summary()
# model.compile('adam', 'categorical_crossentropy', ['accuracy'])
# model.fit(train, epochs=epochs, validation_data=validation)

def unpack(model, training_config, weights):
    restored_model = deserialize(model)
    if training_config is not None:
        restored_model.compile(
            **saving_utils.compile_args_from_training_config(
                training_config
            )
        )
    restored_model.set_weights(weights)
    return restored_model

# Hotfix function
def make_keras_picklable():

    def __reduce__(self):
        model_metadata = saving_utils.model_metadata(self)
        training_config = model_metadata.get("training_config", None)
        model = serialize(self)
        weights = self.get_weights()
        return (unpack, (model, training_config, weights))

    cls = Model
    cls.__reduce__ = __reduce__

# Run the function
make_keras_picklable()


def transform(x):
  x_transform = []
  for i, word_indices in enumerate(x):
    BoW_array = np.zeros((NUM_WORDS,))
    for index in word_indices:
      if index < len(BoW_array):
        BoW_array[index] += 1
    x_transform.append(BoW_array)
  return np.array(x_transform)




class CustomKerasModelWrapper(ModelWrapper):
    def __init__(self, model):
        self.model = model

    def __call__(self, text_input_list):

      x_transform = []
      for i, review in enumerate(text_input_list):
        tokens = [x.strip(",") for x in review.split()]
        BoW_array = np.zeros((NUM_WORDS,))
        for word in tokens:
          if word in vocabulary:
            if vocabulary[word] < len(BoW_array):
              BoW_array[vocabulary[word]] += 1
        x_transform.append(BoW_array)
      x_transform = np.array(x_transform)
      prediction = self.model.predict(x_transform)
      return prediction

model = Sequential()
model.add(Dense(512, activation='relu', input_dim=NUM_WORDS))
model.add(Dropout(0.3))
model.add(Dense(100, activation='relu'))
model.add(Dense(2, activation='sigmoid'))
opt = tf.keras.optimizers.Adam(learning_rate=0.00001)

model.compile(
 optimizer = opt,
loss = "binary_crossentropy",
 metrics = ["accuracy"]
)



(x_train_tokens, y_train), (x_test_tokens, y_test) = tf.keras.datasets.imdb.load_data(
    path="imdb.npz",
    num_words=NUM_WORDS,
    skip_top=0,
    maxlen=None,
    seed=113,
    start_char=1,
    oov_char=2,
    index_from=3
)


index = int(0.9 * len(x_train_tokens))
x_train = transform(x_train_tokens)[:index]
x_test = transform(x_test_tokens)[index:]
y_train = np.array(y_train[:index])
y_test = np.array(y_test[index:])
y_train = to_categorical(y_train)
y_test = to_categorical(y_test)

vocabulary = tf.keras.datasets.imdb.get_word_index(
    path='imdb_word_index.json'
)

results = model.fit(
 x_train, y_train,
 epochs= 1,
 batch_size = 512,
 validation_data = (x_test, y_test)
)


if __name__ == '__main__':
    torch.multiprocessing.freeze_support()


    model_wrapper = CustomKerasModelWrapper(model)
    dataset = HuggingFaceDataset("rotten_tomatoes", None, "test", shuffle=True)

    attack = PWWSRen2019.build(model_wrapper)

    attack_args = AttackArgs(num_examples=10, checkpoint_dir="checkpoints", parallel=True, num_workers_per_device=2)

    attacker = Attacker(attack, dataset, attack_args)

    attacker.attack_dataset()

[qiyanjun]

Great! Possible to add a tutorial in our doc?

…

Sent from my iPhone

On Aug 12, 2021, at 16:32, Sanchit Sinha ***@***.***> wrote:  So I was able to get this to work. The main problem lies with the tf.keras object. You have to nest it inside the if __name__ == '__main__'. Here's the code: import tensorflow as tf import keras import pickle import numpy as np from textattack.models.wrappers import ModelWrapper from textattack.datasets import HuggingFaceDataset from textattack.attack_recipes import PWWSRen2019 import numpy as np # from keras.utils import to_categorical # from tensorflow.keras.optimizers import Adam from tensorflow.keras.utils import to_categorical from tensorflow.keras import models from tensorflow.keras import layers from tensorflow.keras.models import Sequential from tensorflow.keras.layers import Dense from tensorflow.keras.layers import Flatten from tensorflow.keras.layers import Dropout from tensorflow.keras.models import Sequential, load_model, save_model, Model from nltk.tokenize import word_tokenize, RegexpTokenizer from textattack import AttackArgs from textattack.datasets import Dataset from textattack import Attacker import tempfile from tensorflow.python.keras.layers import deserialize, serialize from tensorflow.python.keras.saving import saving_utils import torch NUM_WORDS = 1000 # model = Sequential() # model.add(Embedding(10000, 128, input_length=500, name='embedding')) # model.add(GlobalMaxPooling1D()) # model.add(Dense(n_choice, 'softmax')) # model.summary() # model.compile('adam', 'categorical_crossentropy', ['accuracy']) # model.fit(train, epochs=epochs, validation_data=validation) def unpack(model, training_config, weights): restored_model = deserialize(model) if training_config is not None: restored_model.compile( **saving_utils.compile_args_from_training_config( training_config ) ) restored_model.set_weights(weights) return restored_model # Hotfix function def make_keras_picklable(): def __reduce__(self): model_metadata = saving_utils.model_metadata(self) training_config = model_metadata.get("training_config", None) model = serialize(self) weights = self.get_weights() return (unpack, (model, training_config, weights)) cls = Model cls.__reduce__ = __reduce__ # Run the function make_keras_picklable() def transform(x): x_transform = [] for i, word_indices in enumerate(x): BoW_array = np.zeros((NUM_WORDS,)) for index in word_indices: if index < len(BoW_array): BoW_array[index] += 1 x_transform.append(BoW_array) return np.array(x_transform) class CustomKerasModelWrapper(ModelWrapper): def __init__(self, model): self.model = model def __call__(self, text_input_list): x_transform = [] for i, review in enumerate(text_input_list): tokens = [x.strip(",") for x in review.split()] BoW_array = np.zeros((NUM_WORDS,)) for word in tokens: if word in vocabulary: if vocabulary[word] < len(BoW_array): BoW_array[vocabulary[word]] += 1 x_transform.append(BoW_array) x_transform = np.array(x_transform) prediction = self.model.predict(x_transform) return prediction model = Sequential() model.add(Dense(512, activation='relu', input_dim=NUM_WORDS)) model.add(Dropout(0.3)) model.add(Dense(100, activation='relu')) model.add(Dense(2, activation='sigmoid')) opt = tf.keras.optimizers.Adam(learning_rate=0.00001) model.compile( optimizer = opt, loss = "binary_crossentropy", metrics = ["accuracy"] ) (x_train_tokens, y_train), (x_test_tokens, y_test) = tf.keras.datasets.imdb.load_data( path="imdb.npz", num_words=NUM_WORDS, skip_top=0, maxlen=None, seed=113, start_char=1, oov_char=2, index_from=3 ) index = int(0.9 * len(x_train_tokens)) x_train = transform(x_train_tokens)[:index] x_test = transform(x_test_tokens)[index:] y_train = np.array(y_train[:index]) y_test = np.array(y_test[index:]) y_train = to_categorical(y_train) y_test = to_categorical(y_test) vocabulary = tf.keras.datasets.imdb.get_word_index( path='imdb_word_index.json' ) results = model.fit( x_train, y_train, epochs= 1, batch_size = 512, validation_data = (x_test, y_test) ) if __name__ == '__main__': torch.multiprocessing.freeze_support() model_wrapper = CustomKerasModelWrapper(model) dataset = HuggingFaceDataset("rotten_tomatoes", None, "test", shuffle=True) attack = PWWSRen2019.build(model_wrapper) attack_args = AttackArgs(num_examples=10, checkpoint_dir="checkpoints", parallel=True, num_workers_per_device=2) attacker = Attacker(attack, dataset, attack_args) attacker.attack_dataset() — You are receiving this because you commented. Reply to this email directly, view it on GitHub, or unsubscribe. Triage notifications on the go with GitHub Mobile for iOS or Android.

[XiYuan68]

Thank You so much! It works and makes attacks much faster! BTW for my CNN networks there is a OOM error:

tensorflow.python.framework.errors_impl.InternalError: CUDA runtime implicit initialization on GPU:0 failed. Status: out of memory

and a little workaround can fix it:

gpus = tf.config.experimental.list_physical_devices('GPU')
for gpu in gpus:
  tf.config.experimental.set_memory_growth(gpu, True)

[sanchit97]

@qiyanjun I am working on that right now but there are several problems. There are a lot of namespace issues if we do that inside a notebook because there are different namespaces shared by iPython and multiprocessing. I think a better, easier, and safer way would be to add it as a script. Torch multiprocessing expects all the methods to be in the imports with global scope. I will add a PR detailing changes.