import torch
import torchvision
import torchvision.transforms as transforms
import torch.nn as nn
import torch.optim as optim

# 기본 모델 로드
base_model = torchvision.models.vqa_resnet_finetune(pretrained=True)

# OK-VQA 데이터셋 로드 및 전처리
# (여기에서는 데이터를 로드하는 코드와 전처리 과정을 간략하게 표현합니다)

train_dataset = OKVQADataset('train_data.json', transform=transforms.Compose([transforms.Resize((224, 224)), transforms.ToTensor()]))
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=32, shuffle=True)

# 새로운 레이어 추가
num_classes = len(train_dataset.classes)  # 예시에서는 데이터셋의 클래스 수를 사용
base_model.fc = nn.Linear(base_model.fc.in_features, num_classes)

# GPU 사용 가능 시 GPU로 모델 이동
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
base_model = base_model.to(device)

# Loss 및 Optimizer 정의
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(base_model.parameters(), lr=0.001)

# Fine-tuning
num_epochs = 10
for epoch in range(num_epochs):
    for inputs, labels in train_loader:
        inputs, labels = inputs.to(device), labels.to(device)

        optimizer.zero_grad()
        outputs = base_model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

    print(f'Epoch {epoch+1}/{num_epochs}, Loss: {loss.item()}')
torch.save(base_model.state_dict(), 'git-vqa-finetuned-on-ok-vqa.pth')