lucidrains / vit-pytorch

Implementation of Vision Transformer, a simple way to achieve SOTA in vision classification with only a single transformer encoder, in Pytorch

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

Question about example notebook

repoholder23 opened this issue · comments

Hi, I have a question about the example notebook: is the ViT model pretrained or do we start training from scratch?

In case we start training from scratch, is there a way we can start with a pretrained model instead? Thanks!

The example notebook is training the model from scratch here:

    for data, label in tqdm(train_loader):
        data = data.to(device)
        label = label.to(device)

        output = model(data)
        loss = criterion(output, label)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()