18 lines
509 B
Python
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"]
|
||
|
|
|