Spaces:
Runtime error
Runtime error
File size: 3,597 Bytes
bb5cd12 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 |
import torch
from tqdm import tqdm
from .utils import reduce_losses, to_cuda_half
from torchvision.utils import make_grid
def train_step(config, train_loader, model_engine):
losses = []
for _ in range(config.gradient_accumulation_steps):
images, captions = next(train_loader)
images, captions = images.half().cuda(), captions.cuda()
if config.run_blind:
images = torch.zeros_like(images)
outputs = model_engine(images, captions)
loss = outputs.loss
losses.append(loss)
model_engine.backward(loss)
model_engine.step()
return reduce_losses(torch.mean(torch.stack(losses))).item()
def train_step_classification(config, train_loader, model_engine, return_accuracy=True):
losses = []
if return_accuracy:
accuracies = []
for _ in range(config.gradient_accumulation_steps):
images, captions, class_labels = next(train_loader)
images, captions, class_labels = to_cuda_half(images, captions, class_labels)
if config.run_blind:
images = torch.zeros_like(images)
loss, logits = model_engine(images, captions, class_labels)
losses.append(loss)
if return_accuracy:
argmax_pred = logits.argmax(dim=-1)
accuracies.append((argmax_pred == class_labels).float().mean())
model_engine.backward(loss)
model_engine.step()
loss_reduced = reduce_losses(torch.mean(torch.stack(losses))).item()
if return_accuracy:
accuracy_reduced = reduce_losses(torch.mean(torch.stack(accuracies))).item()
return loss_reduced, accuracy_reduced
return loss_reduced
def eval_step(config, eval_loader, model_engine):
losses = []
for i in tqdm(range(config.eval_steps), "evaluating..."):
images, captions = next(eval_loader)
images, captions = images.half().cuda(), captions.cuda()
if config.run_blind:
images = torch.zeros_like(images)
outputs = model_engine(images, captions)
loss = outputs.loss
losses.append(loss)
return reduce_losses(torch.mean(torch.stack(losses))).item()
def eval_step_classification(config, train_loader, model_engine, return_accuracy=True):
losses = []
if return_accuracy:
accuracies = []
for _ in range(config.gradient_accumulation_steps):
images, captions, class_labels = next(train_loader)
images, captions, class_labels = to_cuda_half(images, captions, class_labels)
if config.run_blind:
images = torch.zeros_like(images)
loss, logits = model_engine(images, captions, class_labels)
losses.append(loss)
if return_accuracy:
argmax_pred = logits.argmax(dim=-1)
accuracies.append((argmax_pred == class_labels).float().mean())
loss_reduced = reduce_losses(torch.mean(torch.stack(losses))).item()
if return_accuracy:
accuracy_reduced = reduce_losses(torch.mean(torch.stack(accuracies))).item()
return loss_reduced, accuracy_reduced
return loss_reduced
def inference_step(config, eval_loader, model_engine):
images, _ = next(eval_loader)
images = images.half().cuda()
if config.run_blind:
images = torch.zeros_like(images)
captions = model_engine(
images, captions=None, inference=True
) # [caption1, caption2, ... b]
width = min(2, images.shape[0])
image_grid = make_grid(images[:width])
caption = ""
for i in range(width):
caption += f"Caption {i}: \n{captions[i]}\n"
return image_grid, caption
|