|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
import torch |
|
from tqdm import tqdm |
|
|
|
from data import dataloader |
|
from model import model, tokenizer, optimizer, load_image |
|
import config |
|
|
|
def get_loss(model, input): |
|
ids = tokenizer(input['text'], return_tensors='pt', padding=True, truncation=True, max_length=config.max_tokens).to(config.device) |
|
|
|
pixel_values = input['image'].to(config.device, config.dtype) |
|
pixel_values = torch.nn.functional.interpolate(pixel_values, (224, 224)) |
|
with torch.cuda.amp.autocast(enabled=True, dtype=config.dtype): |
|
output = model(**ids, labels=ids.input_ids, pixel_values=pixel_values) |
|
|
|
return output.loss |
|
|
|
scaler = torch.cuda.amp.GradScaler() |
|
|
|
for epoch in range(config.epochs): |
|
for ind, sample in tqdm(enumerate(iter(dataloader))): |
|
if sample is None: |
|
continue |
|
|
|
if ind % 100 == 0: |
|
with torch.cuda.amp.autocast(enabled=True, dtype=config.dtype): |
|
response = model.chat(tokenizer=tokenizer, |
|
pixel_values=torch.nn.functional.interpolate( |
|
load_image('/home/ryn_mote/Downloads/horse_style.png').to(config.device, config.dtype), |
|
(224, 224)), |
|
question='<image>\n ', |
|
generation_config = dict(max_new_tokens=config.max_tokens, do_sample=True)) |
|
print('\n\n\n', response, '\n\n\n' ) |
|
|
|
response = model.chat(tokenizer=tokenizer, |
|
pixel_values=torch.nn.functional.interpolate( |
|
load_image('/home/ryn_mote/Downloads/1200px-Andrzej_Person_Kancelaria_Senatu.jpg').to(config.device, config.dtype), |
|
(224, 224)), |
|
question='<image>\n ', |
|
generation_config = dict(max_new_tokens=config.max_tokens, do_sample=True)) |
|
print('\n\n\n', response, '\n\n\n' ) |
|
|
|
loss = get_loss(model, sample) |
|
print(loss.item()) |
|
|
|
scaler.scale(loss).backward() |
|
|
|
optimizer.step() |
|
optimizer.zero_grad() |
|
|
|
|
|
if ind % 1000 == 0: |
|
model.save_pretrained(config.save_path, from_pt=True) |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
