| import os |
| import torch |
| import torch.nn as nn |
| import torch.optim as optim |
| from torchvision import datasets, transforms, models |
| from torch.utils.data import DataLoader |
| from torch.multiprocessing import Process, set_start_method |
|
|
| try: |
| set_start_method('spawn') |
| except RuntimeError: |
| pass |
|
|
| |
| BATCH_SIZE = 32 |
| EPOCHS = 10 |
| NUM_CLASSES = 2 |
| DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu' |
| DATA_ROOTS = [ |
| '/home/ubuntu/vnet/TaoST/Data10kKaggle1', |
| '/home/ubuntu/vnet/TaoST/Data10kKaggle2' |
| ] |
| MODEL_PATHS = [ |
| '/home/ubuntu/vnet/FL/efficientnet_b0_kaggle1.pth', |
| '/home/ubuntu/vnet/FL/efficientnet_b0_kaggle2.pth' |
| ] |
|
|
| def get_loaders(data_root): |
| train_transform = transforms.Compose([ |
| transforms.Resize((224, 224)), |
| transforms.RandomHorizontalFlip(), |
| transforms.ToTensor(), |
| ]) |
| test_transform = transforms.Compose([ |
| transforms.Resize((224, 224)), |
| transforms.ToTensor(), |
| ]) |
| train_set = datasets.ImageFolder(os.path.join(data_root, 'train'), transform=train_transform) |
| test_set = datasets.ImageFolder(os.path.join(data_root, 'test'), transform=test_transform) |
| train_loader = DataLoader(train_set, batch_size=BATCH_SIZE, shuffle=True, num_workers=4) |
| test_loader = DataLoader(test_set, batch_size=BATCH_SIZE, shuffle=False, num_workers=4) |
| return train_loader, test_loader |
|
|
| def train_model(data_root, model_path): |
| train_loader, test_loader = get_loaders(data_root) |
| model = models.efficientnet_b0(weights='IMAGENET1K_V1') |
| model.classifier[1] = nn.Linear(model.classifier[1].in_features, NUM_CLASSES) |
| model = model.to(DEVICE) |
| criterion = nn.CrossEntropyLoss() |
| optimizer = optim.Adam(model.parameters(), lr=1e-4) |
|
|
| for epoch in range(EPOCHS): |
| model.train() |
| running_loss = 0.0 |
| for imgs, labels in train_loader: |
| imgs, labels = imgs.to(DEVICE), labels.to(DEVICE) |
| optimizer.zero_grad() |
| outputs = model(imgs) |
| loss = criterion(outputs, labels) |
| loss.backward() |
| optimizer.step() |
| running_loss += loss.item() * imgs.size(0) |
| print(f"[{data_root}] Epoch {epoch+1}/{EPOCHS}, Loss: {running_loss/len(train_loader.dataset):.4f}") |
|
|
| torch.save(model.state_dict(), model_path) |
| print(f"Saved model to {model_path}") |
|
|
| def main(): |
| p1 = Process(target=train_model, args=(DATA_ROOTS[0], MODEL_PATHS[0])) |
| p2 = Process(target=train_model, args=(DATA_ROOTS[1], MODEL_PATHS[1])) |
| p1.start() |
| p2.start() |
| p1.join() |
| p2.join() |
|
|
| if __name__ == "__main__": |
| main() |
