Introduction:

Neural network models built with TensorFlow.

Installation:

Download models from https://github.com/NoteDance/models and then unzip it to the site-packages folder of your Python environment.

Train:

frommodels.ViTimportViTvit=ViT(
image_size=224,
patch_size=16,
num_classes=1000,
dim=768,
depth=12,
heads=12,
mlp_dim=3072,
pool='cls',
channels=3,
dim_head=64,
drop_rate=0.1,
emb_dropout=0.1
)
loss_fn=tf.keras.losses.SparseCategoricalCrossentropy()
model.compile(optimizer='adam',loss=loss_fn)
model.fit(x_train, y_train, epochs=5)

Distributed training:

frommodels.ViTimportViTstrategy=tf.distribute.MirroredStrategy()
BATCH_SIZE_PER_REPLICA=64GLOBAL_BATCH_SIZE=BATCH_SIZE_PER_REPLICA*strategy.num_replicas_in_syncEPOCHS=10train_dataset=tf.data.Dataset.from_tensor_slices((train_images, train_labels)).shuffle(BUFFER_SIZE).batch(GLOBAL_BATCH_SIZE)
train_dist_dataset=strategy.experimental_distribute_dataset(train_dataset)
withstrategy.scope():
loss_object=tf.keras.losses.SparseCategoricalCrossentropy(
reduction=tf.keras.losses.Reduction.NONE)
defcompute_loss(labels, predictions):
per_example_loss=loss_object(labels, predictions)
returntf.nn.compute_average_loss(per_example_loss, global_batch_size=GLOBAL_BATCH_SIZE)
withstrategy.scope():
vit=ViT(
image_size=224,
patch_size=16,
num_classes=1000,
dim=768,
depth=12,
heads=12,
mlp_dim=3072,
pool='cls',
channels=3,
dim_head=64,
drop_rate=0.1,
emb_dropout=0.1
)
optimizer=tf.keras.optimizers.Adam()
deftrain_step(inputs):
images, labels=inputswithtf.GradientTape() astape:
predictions=vit(images)
loss=compute_loss(labels, predictions)
gradients=tape.gradient(loss, vit.weights)
optimizer.apply_gradients(zip(gradients, vit.weights))
returnloss@tf.function(jit_compile=True)defdistributed_train_step(dataset_inputs):
per_replica_losses=strategy.run(train_step, args=(dataset_inputs,))
returnstrategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses,
axis=None)
forepochinrange(EPOCHS):
total_loss=0.0num_batches=0forxintrain_dist_dataset:
total_loss+=distributed_train_step(x)
num_batches+=1train_loss=total_loss/num_batchestemplate= ("Epoch {}, Loss: {}")
print(template.format(epoch+1, train_loss)

Build models:

Here are some examples of building various neural networks, all in a similar way.

CLIP_large:

frommodels.CLIPimportCLIPclip=CLIP(
embed_dim=1024,
image_resolution=224,
vision_layers=14,
vision_width=1024,
vision_patch_size=32,
context_length=77,
vocab_size=49408,
transformer_width=512,
transformer_heads=8,
transformer_layers=12
)

DiT_B_4:

frommodels.DiTimportDiT_B_4dit=DiT_B_4()

Llama2_7B:

frommodels.Llama2importLlama2llama=Llama2()

ViT

frommodels.ViTimportViTvit=ViT(
image_size=224,
patch_size=16,
num_classes=1000,
dim=768,
depth=12,
heads=12,
mlp_dim=3072,
pool='cls',
channels=3,
dim_head=64,
drop_rate=0.1,
emb_dropout=0.1
)

Assign the trained parameters to the model:

The assign_param function allows you to assign trained parameters, such as downloaded pre-trained parameters, to the parameters of a neural network. These parameters should be stored in a list.

frommodels.assign_paramimportassign_paramassign_param(model.weights,param)