Create app.py

app.py

@@ -0,0 +1,41 @@
+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')