KristofGaming39 commited on
Commit
f3a6572
1 Parent(s): 83a5b01

Create codeimadefortraining.py

Browse files

This is the code that i made for the AI training :>

Files changed (1) hide show
  1. codeimadefortraining.py +60 -0
codeimadefortraining.py ADDED
@@ -0,0 +1,60 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # !!!! MODEL TRAINED BY ME (@KristofGaming39)!!!!
2
+
3
+ import torch
4
+ import torchvision
5
+ from torchvision.models import resnet50
6
+ from torchvision.transforms import transforms
7
+ from torch.utils.data import DataLoader, Dataset
8
+ from PIL import Image
9
+
10
+ class RobloxDataset(Dataset):
11
+ def __init__(self, root_dir, transform=None):
12
+ self.root_dir = root_dir
13
+ self.transform = transform
14
+
15
+ def __len__(self):
16
+ return 15
17
+
18
+ def __getitem__(self, idx):
19
+ img_path = f'{self.root_dir}/Character/human{str(idx+1).zfill(2)}.png'
20
+ image = Image.open(img_path).convert('RGB')
21
+
22
+ if self.transform:
23
+ image = self.transform(image)
24
+
25
+ return image
26
+
27
+ data_transform = transforms.Compose([
28
+ transforms.Resize((224, 224)),
29
+ transforms.ToTensor(),
30
+ transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
31
+ ])
32
+
33
+ dataset = RobloxDataset('/content/dataset', transform=data_transform)
34
+
35
+ data_loader = DataLoader(dataset, batch_size=1, shuffle=True)
36
+
37
+ model = resnet50(pretrained=True)
38
+ model.fc = torch.nn.Linear(in_features=2048, out_features=1)
39
+
40
+ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
41
+ model.to(device)
42
+
43
+ criterion = torch.nn.BCEWithLogitsLoss()
44
+ optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
45
+
46
+ num_epochs = 10
47
+
48
+ for epoch in range(num_epochs):
49
+ for images in data_loader:
50
+ images = images.to(device)
51
+ labels = torch.ones((images.size(0), 1)).to(device)
52
+
53
+ outputs = model(images)
54
+ loss = criterion(outputs, labels)
55
+
56
+ optimizer.zero_grad()
57
+ loss.backward()
58
+ optimizer.step()
59
+
60
+ print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item()}')