This commit is contained in:
Jakub Henyk 2023-05-07 18:29:20 +02:00
parent a0dc9ef9aa
commit 4660aed907
2 changed files with 5 additions and 1 deletions

View File

@ -13,6 +13,7 @@ RUN python3 -m pip install pandas
RUN python3 -m pip install numpy RUN python3 -m pip install numpy
RUN python3 -m pip install torch RUN python3 -m pip install torch
RUN python3 -m pip install torchvision RUN python3 -m pip install torchvision
RUN python3 -m pip install regex
COPY ./zadanie1.py ./ COPY ./zadanie1.py ./
COPY ./Customers.csv ./ COPY ./Customers.csv ./

View File

@ -10,6 +10,7 @@ import torch.nn as nn
import torch.nn.functional as F import torch.nn.functional as F
import torch.optim as optim import torch.optim as optim
import sys import sys
import re
class Net(nn.Module): class Net(nn.Module):
def __init__(self): def __init__(self):
@ -76,7 +77,9 @@ if __name__ == '__main__':
epochs = sys.argv[1] epochs = sys.argv[1]
trainNet(trainloader, criterion, optimizer, int(float(epochs))) epochs_m = re.findall(r'\d+\.\d+', epochs)
trainNet(trainloader, criterion, optimizer, int(float(epochs_m[0])))
PATH = './cifar_net.pth' PATH = './cifar_net.pth'
torch.save(net.state_dict(), PATH) torch.save(net.state_dict(), PATH)