25 lines
2.1 KiB
Text
Executable file
25 lines
2.1 KiB
Text
Executable file
\begin{Verbatim}[commandchars=\\\{\}]
|
|
\PYG{n}{net} \PYG{o}{=} \PYG{n}{Net}\PYG{p}{()}
|
|
|
|
\PYG{n}{device} \PYG{o}{=} \PYG{n}{torch}\PYG{o}{.}\PYG{n}{device}\PYG{p}{(}\PYG{l+s+s2}{\PYGZdq{}cuda:0\PYGZdq{}} \PYG{k}{if} \PYG{n}{torch}\PYG{o}{.}\PYG{n}{cuda}\PYG{o}{.}\PYG{n}{is\PYGZus{}available}\PYG{p}{()} \PYG{k}{else} \PYG{l+s+s2}{\PYGZdq{}cpu\PYGZdq{}}\PYG{p}{)}
|
|
\PYG{n+nb}{print}\PYG{p}{(}\PYG{l+s+s1}{\PYGZsq{}runnning on }\PYG{l+s+si}{\PYGZpc{}s}\PYG{l+s+s1}{\PYGZsq{}} \PYG{o}{\PYGZpc{}} \PYG{n}{device}\PYG{p}{)}
|
|
\PYG{n}{net} \PYG{o}{=} \PYG{n}{net}\PYG{o}{.}\PYG{n}{to}\PYG{p}{(}\PYG{n}{device}\PYG{p}{)}
|
|
|
|
\PYG{n}{loss\PYGZus{}function} \PYG{o}{=} \PYG{n}{nn}\PYG{o}{.}\PYG{n}{CrossEntropyLoss}\PYG{p}{()}
|
|
\PYG{n}{optimizer} \PYG{o}{=} \PYG{n}{optim}\PYG{o}{.}\PYG{n}{Adam}\PYG{p}{(}\PYG{n}{net}\PYG{o}{.}\PYG{n}{parameters}\PYG{p}{(),} \PYG{n}{lr}\PYG{o}{=}\PYG{n}{LEARNING\PYGZus{}RATE}\PYG{p}{)}
|
|
|
|
\PYG{k}{for} \PYG{n}{epoch} \PYG{o+ow}{in} \PYG{n+nb}{range}\PYG{p}{(}\PYG{n}{EPOCHS}\PYG{p}{):}
|
|
\PYG{k}{for} \PYG{n}{data} \PYG{o+ow}{in} \PYG{n}{trainset}\PYG{p}{:}
|
|
\PYG{n}{X}\PYG{p}{,} \PYG{n}{y} \PYG{o}{=} \PYG{n}{data}
|
|
\PYG{n}{net}\PYG{o}{.}\PYG{n}{zero\PYGZus{}grad}\PYG{p}{()}
|
|
\PYG{n}{X} \PYG{o}{=} \PYG{n}{X}\PYG{o}{.}\PYG{n}{to}\PYG{p}{(}\PYG{n}{device}\PYG{p}{)}
|
|
\PYG{n}{output} \PYG{o}{=} \PYG{n}{net}\PYG{p}{(}\PYG{n}{X}\PYG{o}{.}\PYG{n}{view}\PYG{p}{(}\PYG{o}{\PYGZhy{}}\PYG{l+m+mi}{1}\PYG{p}{,} \PYG{n}{n} \PYG{o}{*} \PYG{n}{n}\PYG{p}{))}
|
|
\PYG{n}{output} \PYG{o}{=} \PYG{n}{output}\PYG{o}{.}\PYG{n}{cpu}\PYG{p}{()}
|
|
\PYG{n}{loss} \PYG{o}{=} \PYG{n}{loss\PYGZus{}function}\PYG{p}{(}\PYG{n}{output}\PYG{p}{,} \PYG{n}{y}\PYG{p}{)}
|
|
\PYG{n}{loss}\PYG{o}{.}\PYG{n}{backward}\PYG{p}{()}
|
|
\PYG{n}{optimizer}\PYG{o}{.}\PYG{n}{step}\PYG{p}{()}
|
|
|
|
\PYG{n}{net} \PYG{o}{=} \PYG{n}{net}\PYG{o}{.}\PYG{n}{cpu}\PYG{p}{()}
|
|
\PYG{n}{torch}\PYG{o}{.}\PYG{n}{save}\PYG{p}{(}\PYG{n}{net}\PYG{p}{,} \PYG{l+s+s1}{\PYGZsq{}./nets/net\PYGZus{}\PYGZsq{}} \PYG{o}{+} \PYG{n+nb}{str}\PYG{p}{(}\PYG{n}{epoch}\PYG{p}{)} \PYG{o}{+} \PYG{l+s+s2}{\PYGZdq{}.pt\PYGZdq{}}\PYG{p}{)}
|
|
\PYG{n}{net} \PYG{o}{=} \PYG{n}{net}\PYG{o}{.}\PYG{n}{to}\PYG{p}{(}\PYG{n}{device}\PYG{p}{)}
|
|
\end{Verbatim}
|