Spaces:
Sleeping
Sleeping
File size: 4,151 Bytes
d32e728 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 | import argparse
import torch
import tensorflow as tf
from torch.utils.data import DataLoader, random_split
from torchvision import datasets, transforms
from models.cnn_pytorch import CNN
from models.cnn_tensorflow import build_model
from utils.prep import CLASSES
def get_data_loaders(data_path, batch_size=32):
transform = transforms.Compose([
transforms.Resize((150, 150)),
transforms.ToTensor()
])
train_data = datasets.ImageFolder(
f"{data_path}/seg_train/seg_train",
transform=transform
)
test_data = datasets.ImageFolder(
f"{data_path}/seg_test/seg_test",
transform=transform
)
val_size = int(0.2 * len(train_data))
train_size = len(train_data) - val_size
train_data, val_data = random_split(train_data, [train_size, val_size])
train_loader = DataLoader(train_data, batch_size=batch_size, shuffle=True)
val_loader = DataLoader(val_data, batch_size=batch_size, shuffle=False)
test_loader = DataLoader(test_data, batch_size=batch_size, shuffle=False)
return train_loader, val_loader, test_loader
def train_pytorch(model, train_loader, val_loader, epochs, device):
model.to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.001, weight_decay=1e-4)
criterion = torch.nn.CrossEntropyLoss()
best_loss = float("inf")
for epoch in range(epochs):
# TRAIN
model.train()
total_loss, correct, total = 0, 0, 0
for x, y in train_loader:
x, y = x.to(device), y.to(device)
optimizer.zero_grad()
outputs = model(x)
loss = criterion(outputs, y)
loss.backward()
optimizer.step()
total_loss += loss.item() * y.size(0)
preds = outputs.argmax(1)
correct += (preds == y).sum().item()
total += y.size(0)
train_acc = 100 * correct / total
train_loss = total_loss / total
# VALIDATION
model.eval()
val_loss, val_correct, val_total = 0, 0, 0
with torch.no_grad():
for x, y in val_loader:
x, y = x.to(device), y.to(device)
outputs = model(x)
loss = criterion(outputs, y)
val_loss += loss.item() * y.size(0)
preds = outputs.argmax(1)
val_correct += (preds == y).sum().item()
val_total += y.size(0)
val_acc = 100 * val_correct / val_total
val_loss = val_loss / val_total
print(f"Epoch {epoch+1}/{epochs} | "
f"Train Loss {train_loss:.4f} Acc {train_acc:.2f}% | "
f"Val Loss {val_loss:.4f} Acc {val_acc:.2f}%")
# save best model
if val_loss < best_loss:
best_loss = val_loss
torch.save(model.state_dict(), "pytorch_model.pth")
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--model", required=True, choices=["pytorch", "tensorflow"])
parser.add_argument("--epochs", type=int, default=25)
parser.add_argument("--data", type=str, required=True)
args = parser.parse_args()
device = "cuda" if torch.cuda.is_available() else "cpu"
# PYTORCH
if args.model == "pytorch":
train_loader, val_loader, test_loader = get_data_loaders(args.data)
model = CNN(num_classes=len(CLASSES))
train_pytorch(model, train_loader, val_loader, args.epochs, device)
# TENSORFLOW
else:
train_ds = tf.keras.preprocessing.image_dataset_from_directory(
f"{args.data}/seg_train/seg_train",
image_size=(150, 150),
batch_size=32
)
val_ds = tf.keras.preprocessing.image_dataset_from_directory(
f"{args.data}/seg_test/seg_test",
image_size=(150, 150),
batch_size=32
)
model = build_model(num_classes=len(CLASSES))
model.fit(
train_ds,
validation_data=val_ds,
epochs=args.epochs
)
model.save("tensorflow_model.keras")
if __name__ == "__main__":
main() |