`# Define the custom dataset
class CustomDataset(Dataset):
def init(self, csv_file, root_dir, transform=None):
self.annotations = pd.read_csv(csv_file)
self.root_dir = root_dir
self.transform = transform
def __len__(self):
return len(self.annotations)
def __getitem__(self, index):
img_path = os.path.join(self.root_dir, f"image_{index}.jpg")
image = Image.open(img_path).convert("RGB")
y_label = self.annotations.iloc[index, 2] # Assuming label is in the third column
if self.transform:
image = self.transform(image)
return image, y_label
Define transformations
transform = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
Hyperparameters
batch_size = 32
learning_rate = 0.001
num_epochs = 25
num_classes = 1000 # Adjust based on your dataset
Load data
dataset = CustomDataset(
csv_file='C:\Users\dmitr\AiProjectTinyVit\TinyViT\Images\locations.csv',
root_dir='C:\Users\dmitr\AiProjectTinyVit\TinyViT\Images\images',
transform=transform
)
train_loader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
Initialize model, loss, and optimizer
model = tiny_vit_21m_224(pretrained=True)
model.head = nn.Linear(model.head.in_features, num_classes) # Adjust for your number of classes
model = model.to('cuda' if torch.cuda.is_available() else 'cpu')
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=learning_rate)
Training loop
for epoch in range(num_epochs):
for batch_idx, (data, targets) in enumerate(train_loader):
data = data.to('cuda' if torch.cuda.is_available() else 'cpu')
targets = targets.to('cuda' if torch.cuda.is_available() else 'cpu')
# Forward pass
scores = model(data)
loss = criterion(scores, targets)`