s452627/ium_452627
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

fix13

Browse Source
This commit is contained in:
Jakub Henyk 2023-05-07 18:31:23 +02:00
parent 7130fa7cdd
commit 3e53e29514
1 changed files with 2 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
train.py
View File

@ -10,7 +10,6 @@ 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,10 +75,9 @@ if __name__ == '__main__':
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9) optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)
epochs = sys.argv[1] epochs = sys.argv[1]
print(epochs)
epochs_m = re.findall(r'\d+\.\d+', epochs) trainNet(trainloader, criterion, optimizer, int(float(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)