Introduction:

Installation:

Train:

from models.ViT import ViT
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
)
loss_fn = tf.keras.losses.SparseCategoricalCrossentropy()
model.compile(optimizer='adam',loss=loss_fn)
model.fit(x_train, y_train, epochs=5)

Distributed training:

from models.ViT import ViT
strategy = tf.distribute.MirroredStrategy()
BATCH_SIZE_PER_REPLICA = 64
GLOBAL_BATCH_SIZE = BATCH_SIZE_PER_REPLICA * strategy.num_replicas_in_sync
EPOCHS = 10
train_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)
with strategy.scope():
 loss_object = tf.keras.losses.SparseCategoricalCrossentropy(
 reduction=tf.keras.losses.Reduction.NONE)
 def compute_loss(labels, predictions):
 per_example_loss = loss_object(labels, predictions)
 return tf.nn.compute_average_loss(per_example_loss, global_batch_size=GLOBAL_BATCH_SIZE)
with strategy.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()
def train_step(inputs):
 images, labels = inputs
 with tf.GradientTape() as tape:
 predictions = vit(images)
 loss = compute_loss(labels, predictions)
 gradients = tape.gradient(loss, vit.weights)
 optimizer.apply_gradients(zip(gradients, vit.weights))
 return loss
@tf.function(jit_compile=True)
def distributed_train_step(dataset_inputs):
 per_replica_losses = strategy.run(train_step, args=(dataset_inputs,))
 return strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses,
 axis=None)
for epoch in range(EPOCHS):
 total_loss = 0.0
 num_batches = 0
 for x in train_dist_dataset:
 total_loss += distributed_train_step(x)
 num_batches += 1
 train_loss = total_loss / num_batches
 template = ("Epoch {}, Loss: {}")
 print(template.format(epoch + 1, train_loss)

Build models:

from models.CLIP import CLIP
clip=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
 )

from models.DiT import DiT_B_4
dit=DiT_B_4()

from models.Llama2 import Llama2
llama=Llama2()

from models.ViT import ViT
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
)

Assign the trained parameters to the model:

from models.assign_param import assign_param
assign_param(model.weights,param)