''' 
1. torch.tensor와 torch.as_tensor의 차이
'''
a = torch.tensor(original)
b = torch.as_tensor(original) # 얕은 복사, 값이 바뀐다.

original[...] = 1
print('a: ', a)
print('b: ', b)

# a: tensor([[0., 0., 0.],[0., 0., 0.]], dtype=torch.float64)
# b: tensor([[1., 1., 1.], [1., 1., 1.]], dtype=torch.float64)

''' 
2. network의 shape을 보는 방법 : named_parameter, parameter 함수
'''

linear = nn.Linear(in_features=2, out_features=3)
for name, param in linear.named_parameters():
    print('linear', name, param.size())

conv = nn.Conv2d(in_channels=2, out_channels=3, kernel_size=1)
for name, param in conv.named_parameters():
    print('conv', name, param.size())

# linear weight torch.Size([3, 2])
# linear bias torch.Size([3])
# conv weight torch.Size([3, 2, 1, 1])
# conv bias torch.Size([3])

''' 
3. nn.Sequential, nn.ParameterList, nn.ParameterDict, nn.ModuleList, nn.ModuleDict
'''

class Module(nn.Module):
    def __init__(self):
        super().__init__()
            
        self.network = nn.Sequential(nn.Linear(2, 3), nn.Linear(3, 4))
        
        self.params1 = nn.ParameterList([
            nn.Parameter(torch.zeros(2)),
            nn.Parameter(torch.ones(2))
        ])
        self.params2 = nn.ParameterDict({
            'param1': nn.Parameter(torch.zeros(2)),
            'param2': nn.Parameter(torch.ones(2))
        })
        self.linears1 = nn.ModuleList([nn.Linear(2, 2), nn.Linear(2, 2)])
        self.linears2 = nn.ModuleDict({'linear1': nn.Linear(2, 2), 'linear2': nn.Linear(2, 2)})

    def forward(self, x):
        return self.network(x)

''' 
4. torch.cuda.memory_allocated
'''
print(torch.cuda.memory_allocated())

''' 
5. Autograd & Optimizer
'''
a = torch.ones(2, 3, requires_grad=True)
print(a)
print(a+2)

# tensor([[1., 1., 1.],[1., 1., 1.]], requires_grad=True) 연산을 기록한다는 것이 나온다.
# tensor([[3., 3., 3.],[3., 3., 3.]], grad_fn=<AddBackward0>) 이와 같이 연산을 tracking한다.

a = torch.ones(2, 3, requires_grad=True)
print(a.detach() + 2)

a = torch.ones(2, 3, requires_grad=True)
with torch.no_grad():
    print(a + 2)
# detach와 torch.no_grad는 같은 효과

a = torch.ones(2, 3, requires_grad=True)
b = a + 2
c = (2 * b).sum()
return torch.autograd.grad(c, [a])
# 이렇게 하면 gradient가 return 된다.

return torch.autograd.grad(loss, module.parameters())
# 이렇게 하면 model 안에 있는 gradient가 return 된다

grads = torch.autograd.grad(loss, module.parameters())
for param, grad in zip(module.parameters(), grads):
    param.data = param.data - 0.001*grad.data
# gradient update는 위와 같이 한다.

optimizer.zero_grad() 
loss.backward() # grad를 계산
optimizer.step() # grad update

''' 
6. list를 sequential에 넣는 것
'''
class Classifier(nn.Module):
    def __init__(self):
        super().__init__()
        layers = []
        in_dim = w*h*c
        for dim in hidden_dims:
            out_dim = dim
            layers.append(nn.Linear(in_dim, out_dim))
            layers.append(nn.ReLU())
            in_dim = out_dim
        layers.append(nn.Linear(out_dim, n_class))
        self.network = nn.Sequential(*layers)