from torch. optim import Adam
class TestNet ( nn. Module) :
def __init__ ( self) :
super ( ) . __init__( )
self. hidden = nn. Sequential(
nn. Linear( 13 , 10 ) ,
nn. ReLU( )
)
self. regression = nn. Linear( 10 , 1 )
def forward ( self, x) :
x = self. hidden( x)
output = self. regression( x)
return output
testnet = TestNet( )
print ( testnet)
optimizer = Adam( testnet. parameters( ) , lr= 0.001 )
optimizer = Adam(
[ { "params" : testnet. hidden. parameters( ) , "lr" : 0.0001 } ,
{ "params" : testnet. regression. parameters( ) , "lr" : 0.01 } ] ,
lr= 1e-2
)
定义好优化器后,需要将optimizer.zero_grad()方法和optimizer.step()方法一起使用,对网络中的参数进行更新。
for x, y in dataloader:
optimizer. zero_grad( )
output = testnet. forward( x)
loss = loss_fn( output, target)
loss. backward( )
optimizer. step( )