ARB Security Solutions, LLC.

Transfer Learning, TensorBoardX, And SharePoint Online

Import dat torch yo! This script configures the data transformations for the custom dataset using data_transforms. It defines the ResNet-50 model for transfer learning and replaces the final fully connected layer with a new one suitable for the custom classification task. The script specifies hyperparameters such as the number of classes, number of epochs, batch size, and learning rate. The train_model function is used to train the model, calculate metrics (loss and accuracy), and log them using TensorBoardX. The main script sets up the SharePoint Online connection, configures the device for training, creates the data loaders for the custom dataset, defines the loss function and optimizer, and creates a TensorBoardX writer for logging. Finally, it calls the train_model function to train the model, saves the trained model to a file, and uploads it to the SharePoint Online document library.

import os
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision.models as models
from torchvision import transforms, datasets
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
from shareplum import Site
from shareplum import Office365

# SharePoint Online credentials and site information
username = 'your_username'
password = 'your_password'
site_url = 'https://your_sharepoint_site_url'
document_library = 'Documents'

# Custom dataset
data_transforms = {
    'train': transforms.Compose([
        transforms.RandomResizedCrop(224),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ]),
    'val': transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])
}

# Data directory
data_dir = 'path_to_data'

# Transfer learning model
model = models.resnet50(pretrained=True)
num_ftrs = model.fc.in_features
model.fc = nn.Linear(num_ftrs, num_classes)

# Hyperparameters
num_classes = 10
num_epochs = 10
batch_size = 32
learning_rate = 0.001

# Training function
def train_model(model, dataloaders, criterion, optimizer, device, writer):
    model.train()
    for epoch in range(num_epochs):
        running_loss = 0.0
        correct_predictions = 0

        for inputs, labels in dataloaders['train']:
            inputs, labels = inputs.to(device), labels.to(device)
            optimizer.zero_grad()

            # Forward pass
            outputs = model(inputs)
            _, predictions = torch.max(outputs, 1)
            loss = criterion(outputs, labels)

            # Backward pass and optimization
            loss.backward()
            optimizer.step()

            running_loss += loss.item()
            correct_predictions += torch.sum(predictions == labels.data)

        # Calculate metrics and log to TensorBoardX
        epoch_loss = running_loss / len(dataloaders['train'])
        epoch_accuracy = correct_predictions.double() / len(dataloaders['train'].dataset)

        writer.add_scalar('Loss/train', epoch_loss, epoch)
        writer.add_scalar('Accuracy/train', epoch_accuracy, epoch)

        print(f'Epoch [{epoch + 1}/{num_epochs}], Loss: {epoch_loss:.4f}, Accuracy: {epoch_accuracy:.4f}')

    writer.close()

# Main script
if __name__ == '__main__':
    # SharePoint Online connection
    with Office365(site_url, username=username, password=password) as auth:
        site = Site(site_url, auth=auth)
        folder = site.Folder(f'{document_library}/DeepLearning')

        # Device configuration
        device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

        # Data loaders
        image_datasets = {x: datasets.ImageFolder(os.path.join(data_dir, x), data_transforms[x]) for x in ['train', 'val']}
        dataloaders = {x: DataLoader(imagedatasets[x], batch_size=batch_size, shuffle=True, num_workers=4) for x in ['train', 'val']}

        # Loss function and optimizer
        criterion = nn.CrossEntropyLoss()
        optimizer = optim.SGD(model.parameters(), lr=learning_rate, momentum=0.9)

        # TensorBoardX writer for logging
        writer = SummaryWriter()

        # Training the model
        train_model(model, dataloaders, criterion, optimizer, device, writer)

        # Save the trained model
        model_filename = 'trained_model.pth'
        torch.save(model.state_dict(), model_filename)
        with open(model_filename, 'rb') as model_file:
            folder.upload_file(model_file, os.path.basename(model_filename))

    print("Model and logs uploaded to SharePoint Online.")

Comments are closed.