Create Train.py
Browse files
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()}')
|