Spaces:
Runtime error
Runtime error
#try this with pytorch directly | |
import torch | |
import torch.nn as nn | |
import torch.optim as optim | |
from torchvision import models | |
from torch.utils.data import DataLoader | |
from torchvision.transforms import transforms | |
from PIL import Image as PILImage | |
import pandas as pd | |
import os | |
# Define the path to the train thumbnails | |
trainPath = Path('/kaggle/input/UBC-OCEAN/train_images') | |
# Define the transform to apply to each image | |
transform = transforms.Compose([ | |
transforms.Resize(460), | |
transforms.CenterCrop(224), | |
transforms.ToTensor(), | |
transforms.Normalize(*imagenet_stats) | |
]) | |
# Define a function to load an image from a file | |
def load_image(filename): | |
image = PILImage.open(filename) | |
image = transform(image) | |
return image | |
# Define a function to load the train dataset | |
def load_train_dataset(): | |
# Load the train CSV file | |
df = pd.read_csv("/kaggle/input/UBC-OCEAN/train.csv") | |
# Create a list of tuples containing the image path and label | |
image_list = [(os.path.join(trainPath, str(row['image_id']) + '.png'), row['label']) for index, row in df.iterrows()] | |
# Create a dataset from the list of image tuples | |
dataset = [(load_image(f), l) for f, l in image_list] | |
return dataset | |
# Load the train dataset | |
train_dataset = load_train_dataset() | |
# Create a dataloader for the train dataset | |
train_dataloader = DataLoader(train_dataset, batch_size=16, shuffle=True) | |
# Define the model architecture | |
model = models.resnet18(pretrained=True) | |
num_ftrs = model.fc.in_features | |
model.fc = nn.Linear(num_ftrs, 2) | |
# Move the model to the GPU | |
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") | |
model = model.to(device) | |
# Define the loss function and optimizer | |
criterion = nn.CrossEntropyLoss() | |
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9) | |
# Train the model | |
model.train() | |
for epoch in range(5): | |
running_loss = 0.0 | |
for i, batch in enumerate(train_dataloader): | |
# Get the inputs and labels | |
inputs, labels = batch | |
inputs, labels = inputs.to(device), labels.to(device) | |
# Zero the parameter gradients | |
optimizer.zero_grad() | |
# Forward pass | |
outputs = model(inputs) | |
loss = criterion(outputs, labels) | |
# Backward pass and optimization | |
loss.backward() | |
optimizer.step() | |
# Print statistics | |
running_loss += loss.item() | |
if i % 100 == 99: | |
print('[%d, %5d] loss: %.3f' % (epoch + 1, i + 1, running_loss / 100)) | |
running_loss = 0.0 | |
print('Finished training') | |
# Save the model | |
torch.save(model.state_dict(), 'ovarianCancerModel.pth') |