안녕하세요. 현재 pytorch의 autogradient의 값을 이용해 loss 함수를 정의하려고 합니다.
제가 이해하기로는 pytorch의 경우 autogradient가 각 데이터 샘플 별로 따로 계산되어 하는 줄로 알고 있는데
샘플의 개수가 많을 때 계산시간이 너무 많이 소요됩니다.
혹시 샘플 전체에 대해서 autgradient를 계산할 수 있는 방법이 있을까요? 아니면 아래 제 코드에서 계산 시간을 단축할 수 있는 다른 방법이 있을 지 여쭤봅니다.
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.linear_relu_stack = nn.Sequential(
nn.Linear(2, 32),
nn.ReLU(),
nn.Linear(32, 64),
nn.ReLU(),
nn.Linear(64, 128),
nn.ReLU(),
nn.Linear(128, 64),
nn.ReLU(),
nn.Linear(64, 1),
)
def forward(self, xtpair):
u = self.linear_relu_stack(xtpair)
return u
n = 3000
x = np.linspace(0., np.pi, n)
t = np.linspace(0., 1., n)
ref_ft = lambda xt: np.sin(xt[:,0]) * np.cos(xt[:,1])
##########################################################################
def loss_custom(IN, pred):
nsample = IN.shape[0]
loss_sum = 0.
for i in range(nsample):
first = torch.autograd.grad(outputs=pred[i], inputs=IN, create_graph=True)[0][i] ## first-order derivative = (u_x, u_t)
loss_sum += first[-1:]**2.
return loss_sum
##########################################################################
torch.manual_seed(0)
model = Net()
optimizer = torch.optim.Adam(model.parameters(), lr=5e-3)
##########################################################################
max_ep = 10
for ep in range(max_ep):
model.train()
IN = torch.FloatTensor(np.vstack([x,t]).T)
ref = ref_ft(IN)
IN.requires_grad = True
pred = model(IN)
start_time = time.time()
loss = loss_custom(IN, pred)
print('loss = %.3E (comput time = %.2f (sec))' % (loss.item(), time.time() - start_time))
optimizer.zero_grad()
loss.backward()
optimizer.step()
