<filename>main.py
import torch
import torch.nn as nn
from vivit import ViViT
import numpy as np
import os
import torch.nn.functional as F
from datasets import haa500_instruments
from torch.optim import AdamW
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')


#define hyperparameters
batch_size = 1
eval_freq = 1

# define transform
import video_transforms
scale = 1
input_size = int(224 * scale)
width = int(340 * scale)
height = int(256 * scale)
length=32
clip_mean = [0.5, 0.5, 0.5] * length
clip_std = [0.5, 0.5, 0.5] * length
scale_ratios = [1.0, 0.875, 0.75, 0.66]
normalize = video_transforms.Normalize(mean=clip_mean, std=clip_std)
train_transform = video_transforms.Compose([
                    video_transforms.MultiScaleCrop((input_size, input_size), scale_ratios),
                    video_transforms.RandomHorizontalFlip(),
                    video_transforms.ToTensor(),
                    normalize,
                ])

val_transform = video_transforms.Compose([
        video_transforms.CenterCrop((input_size)),
        video_transforms.ToTensor(),
        normalize,
    ])


#define dataset and dataloader
train_split_file = "./datasets/settings/haa500_instruments/train_rgb_split1.txt"
val_split_file = "./datasets/settings/haa500_instruments/val_rgb_split1.txt"

dataset_name = "haa500_instruments"
dataset_root = "./datasets/haa500_instruments"
train_dataset = haa500_instruments(train_split_file=train_split_file,
                                    dataset_root=dataset_root,
                                    phase="train",
                                    new_length=length,
                                    new_width=width,
                                    new_height=height,
                                    video_transform=train_transform)

val_dataset = haa500_instruments(train_split_file=val_split_file,
                                dataset_root=dataset_root,
                                  phase="val",
                                  new_length=length,
                                  new_width=width,
                                  new_height=height,
                                  video_transform=val_transform)

print('{} samples found, {} train samples and {} test samples.'.format(len(val_dataset)+len(train_dataset),
                                                                       len(train_dataset),
                                                                       len(val_dataset)))

train_loader = torch.utils.data.DataLoader(
    train_dataset,
    batch_size=batch_size, shuffle=True, pin_memory=True)
val_loader = torch.utils.data.DataLoader(
    val_dataset,
    batch_size=batch_size, shuffle=False, pin_memory=True)



def train_one_epoch(train_loader, model, criterion, optimizer):
    model.train()
    loss_total = 0
    total_sample_count = 0
    for i, (inputs, targets) in enumerate(train_loader):
        batch_size = targets.shape[0]
        inputs = inputs.view(-1, 3, length, input_size, input_size).transpose(1, 2).to(device)
        targets = targets.to(device)
        outputs = model(inputs)
        preds = torch.argmax(outputs)
        print(preds)
        loss = criterion(outputs, targets)
        loss_total += loss * batch_size
        total_sample_count += batch_size
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()
    return loss_total / total_sample_count

def validate(val_loader, model):
    model.eval()
    total_samples = 0
    total_correct_counts = 0
    for i, (inputs, targets) in enumerate(val_loader):
        total_samples += targets.shape[0]
        inputs = inputs.view(-1, 3, length, input_size, input_size).transpose(1, 2).to(device)
        targets = targets.to(device)
        outputs = model(inputs)
        preds = torch.argmax(outputs)
        correct_preds_count = sum(preds == targets).item()
        total_correct_counts += correct_preds_count
    top1acc = total_correct_counts/total_samples
    return top1acc

def train(train_loader, model, criterion, optimizer, epochs=200):
    for epoch in range(epochs):
        loss = train_one_epoch(train_loader, model, criterion, optimizer)
        if epoch % eval_freq == 0:
            top1_acc = validate(val_loader, model)
            print("Epoch: {}    Training Loss: {}    Validation Top1 Acc: {}".format(epoch, loss, top1_acc))


# define model
model = ViViT(image_size=224,#
             patch_size=16,#
             tubelet_temporal_size=2,#
             num_classes=51,#
             num_frames=32,#
             dim=768,#
             layer_spacial=12,#
             layer_temporal=4,#
             heads=12,#
             dim_head=64,
             dropout=0.,
             emb_dropout=0,
            mlp_dim=3072,
              pretrain=True)
parameters = filter(lambda p: p.requires_grad, model.parameters())
parameters = sum([np.prod(p.size()) for p in parameters]) / 1_000_000
print('Trainable Parameters: %.3fM' % parameters)

#define optimizer, criterion
criterion = nn.CrossEntropyLoss().to(device)
optimizer = AdamW(model.parameters(), lr=1e-5, weight_decay=1e-3)

train(train_loader, model, criterion, optimizer, 200)


<filename>main.py
import torch
import torch.nn as nn
from vivit import ViViT
import numpy as np
import os
import torch.nn.functional as F
from datasets import haa500_instruments
from torch.optim import AdamW
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')


#define hyperparameters
batch_size = 1
eval_freq = 1

# define transform
import video_transforms
scale = 1
input_size = int(224 * scale)
width = int(340 * scale)
height = int(256 * scale)
length=32
clip_mean = [0.5, 0.5, 0.5] * length
clip_std = [0.5, 0.5, 0.5] * length
scale_ratios = [1.0, 0.875, 0.75, 0.66]
normalize = video_transforms.Normalize(mean=clip_mean, std=clip_std)
train_transform = video_transforms.Compose([
                    video_transforms.MultiScaleCrop((input_size, input_size), scale_ratios),
                    video_transforms.RandomHorizontalFlip(),
                    video_transforms.ToTensor(),
                    normalize,
                ])

val_transform = video_transforms.Compose([
        video_transforms.CenterCrop((input_size)),
        video_transforms.ToTensor(),
        normalize,
    ])


#define dataset and dataloader
train_split_file = "./datasets/settings/haa500_instruments/train_rgb_split1.txt"
val_split_file = "./datasets/settings/haa500_instruments/val_rgb_split1.txt"

dataset_name = "haa500_instruments"
dataset_root = "./datasets/haa500_instruments"
train_dataset = haa500_instruments(train_split_file=train_split_file,
                                    dataset_root=dataset_root,
                                    phase="train",
                                    new_length=length,
                                    new_width=width,
                                    new_height=height,
                                    video_transform=train_transform)

val_dataset = haa500_instruments(train_split_file=val_split_file,
                                dataset_root=dataset_root,
                                  phase="val",
                                  new_length=length,
                                  new_width=width,
                                  new_height=height,
                                  video_transform=val_transform)

print('{} samples found, {} train samples and {} test samples.'.format(len(val_dataset)+len(train_dataset),
                                                                       len(train_dataset),
                                                                       len(val_dataset)))

train_loader = torch.utils.data.DataLoader(
    train_dataset,
    batch_size=batch_size, shuffle=True, pin_memory=True)
val_loader = torch.utils.data.DataLoader(
    val_dataset,
    batch_size=batch_size, shuffle=False, pin_memory=True)



def train_one_epoch(train_loader, model, criterion, optimizer):
    model.train()
    loss_total = 0
    total_sample_count = 0
    for i, (inputs, targets) in enumerate(train_loader):
        batch_size = targets.shape[0]
        inputs = inputs.view(-1, 3, length, input_size, input_size).transpose(1, 2).to(device)
        targets = targets.to(device)
        outputs = model(inputs)
        preds = torch.argmax(outputs)
        print(preds)
        loss = criterion(outputs, targets)
        loss_total += loss * batch_size
        total_sample_count += batch_size
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()
    return loss_total / total_sample_count

def validate(val_loader, model):
    model.eval()
    total_samples = 0
    total_correct_counts = 0
    for i, (inputs, targets) in enumerate(val_loader):
        total_samples += targets.shape[0]
        inputs = inputs.view(-1, 3, length, input_size, input_size).transpose(1, 2).to(device)
        targets = targets.to(device)
        outputs = model(inputs)
        preds = torch.argmax(outputs)
        correct_preds_count = sum(preds == targets).item()
        total_correct_counts += correct_preds_count
    top1acc = total_correct_counts/total_samples
    return top1acc

def train(train_loader, model, criterion, optimizer, epochs=200):
    for epoch in range(epochs):
        loss = train_one_epoch(train_loader, model, criterion, optimizer)
        if epoch % eval_freq == 0:
            top1_acc = validate(val_loader, model)
            print("Epoch: {}    Training Loss: {}    Validation Top1 Acc: {}".format(epoch, loss, top1_acc))


# define model
model = ViViT(image_size=224,#
             patch_size=16,#
             tubelet_temporal_size=2,#
             num_classes=51,#
             num_frames=32,#
             dim=768,#
             layer_spacial=12,#
             layer_temporal=4,#
             heads=12,#
             dim_head=64,
             dropout=0.,
             emb_dropout=0,
            mlp_dim=3072,
              pretrain=True)
parameters = filter(lambda p: p.requires_grad, model.parameters())
parameters = sum([np.prod(p.size()) for p in parameters]) / 1_000_000
print('Trainable Parameters: %.3fM' % parameters)

#define optimizer, criterion
criterion = nn.CrossEntropyLoss().to(device)
optimizer = AdamW(model.parameters(), lr=1e-5, weight_decay=1e-3)

train(train_loader, model, criterion, optimizer, 200)


#define hyperparameters # define transform #define dataset and dataloader # define model # # # # # # # # # #define optimizer, criterion