Inteligentny_Wozek/NeuralNetwork/neural_network_learning.py

18 lines
509 B
Python

from src.torchvision_resize_dataset import combined_dataset
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torchvision import datasets
from torchvision.transforms import Compose, Lambda, ToTensor
import matplotlib.pyplot as plt
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
train_loader = DataLoader(
combined_dataset, #dataset of images
batch_size=256, # accuracy
shuffle=True # rand order
)
classes = ["package", "list"]