Yuzmi commited on
Commit
27010fd
1 Parent(s): 4de61ef

Create Train.py

Browse files
Files changed (1) hide show
  1. Train.py +36 -0
Train.py ADDED
@@ -0,0 +1,36 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ from torchvision import datasets, transforms
3
+ import torch.nn as nn
4
+ import torch.optim as optim
5
+
6
+ # Define your model class
7
+ class TatsukichiHayamaClassifier(nn.Module):
8
+ # ... (your model definition)
9
+
10
+ # Load dataset from PyTorch's ImageFolder
11
+ train_dataset = datasets.ImageFolder(root="TatsukichiHayamaDataset", transform=transforms.ToTensor())
12
+
13
+ # Create a DataLoader for training
14
+ dataloader = torch.utils.data.DataLoader(train_dataset, batch_size=32, shuffle=True)
15
+
16
+ # Create an instance of TatsukichiHayamaClassifier
17
+ your_num_classes = 10 # Adjust this based on your dataset
18
+ model = TatsukichiHayamaClassifier(num_classes=your_num_classes)
19
+
20
+ # Model, criterion, and optimizer
21
+ criterion = nn.CrossEntropyLoss()
22
+ optimizer = optim.Adam(model.parameters(), lr=0.001)
23
+
24
+ # Training loop
25
+ num_epochs = 10
26
+
27
+ for epoch in range(num_epochs):
28
+ model.train()
29
+ for images, labels in dataloader:
30
+ optimizer.zero_grad()
31
+ outputs = model(images)
32
+ loss = criterion(outputs, labels)
33
+ loss.backward()
34
+ optimizer.step()
35
+
36
+ print(f'Epoch {epoch+1}/{num_epochs}, Loss: {loss.item()}')